From 170e7c7bb8ebacdce1171a5cb66da1c825cfe9ab Mon Sep 17 00:00:00 2001
From: tigerwill90 <sylvainmuller@protonmail.com>
Date: Tue, 31 Dec 2024 15:20:58 +0100
Subject: [PATCH] feat: improve warpf and warph function performance

---
 context.go      | 31 +++++++++++++++++++++++++++----
 context_test.go | 19 +++++++++++++++++++
 txn.go          |  5 +++++
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/context.go b/context.go
index 0ae88e4..b942b67 100644
--- a/context.go
+++ b/context.go
@@ -219,14 +219,30 @@ func (c *cTx) Params() iter.Seq[Param] {
 
 // Param retrieve a matching wildcard segment by name.
 func (c *cTx) Param(name string) string {
-	for p := range c.Params() {
-		if p.Key == name {
-			return p.Value
+	if c.tsr {
+		for i := range *c.tsrParams {
+			if (*c.tsrParams)[i].Key == name {
+				return (*c.tsrParams)[i].Value
+			}
+		}
+		return ""
+	}
+
+	for i := range *c.params {
+		if (*c.params)[i].Key == name {
+			return (*c.params)[i].Value
 		}
 	}
 	return ""
 }
 
+func (c *cTx) ParamsLen() int {
+	if c.tsr {
+		return len(*c.tsrParams)
+	}
+	return len(*c.params)
+}
+
 // Path returns the request URL path.
 func (c *cTx) Path() string {
 	return c.req.URL.Path
@@ -402,7 +418,14 @@ func (c *cTx) getQueries() url.Values {
 // The route parameters are being accessed by the wrapped handler through the context.
 func WrapF(f http.HandlerFunc) HandlerFunc {
 	return func(c Context) {
-		var params Params = slices.Collect(c.Params())
+
+		var params Params
+		if ps, ok := c.(interface{ ParamsLen() int }); ok {
+			params = make(Params, 0, ps.ParamsLen())
+			params = slices.AppendSeq(params, c.Params())
+		} else {
+			params = slices.Collect(c.Params())
+		}
 		if len(params) > 0 {
 			ctx := context.WithValue(c.Request().Context(), paramsKey, params)
 			f.ServeHTTP(c.Writer(), c.Request().WithContext(ctx))
diff --git a/context_test.go b/context_test.go
index 1b1c8f4..6ddf940 100644
--- a/context_test.go
+++ b/context_test.go
@@ -513,3 +513,22 @@ func TestWrapH(t *testing.T) {
 		})
 	}
 }
+
+// BenchmarkWrapF-16    	 2211204	       559.8 ns/op	     696 B/op	      10 allocs/op
+func BenchmarkWrapF(b *testing.B) {
+	req := httptest.NewRequest(http.MethodGet, "https://example.com/a/b/c", nil)
+	w := httptest.NewRecorder()
+
+	f := New()
+	f.MustHandle(http.MethodGet, "/{a}/{b}/{c}", WrapF(func(w http.ResponseWriter, r *http.Request) {
+
+	}))
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for range b.N {
+		f.ServeHTTP(w, req)
+	}
+
+}
diff --git a/txn.go b/txn.go
index b3b2ebc..9c84ef9 100644
--- a/txn.go
+++ b/txn.go
@@ -112,6 +112,8 @@ func (txn *Txn) Delete(method, pattern string) error {
 	return nil
 }
 
+// Truncate remove all routes for the provided methods. If not methods are provided,
+// all routes are truncated.
 func (txn *Txn) Truncate(methods ...string) error {
 	if txn.rootTxn == nil {
 		panic(ErrSettledTxn)
@@ -228,6 +230,9 @@ func (txn *Txn) Iter() Iter {
 
 // Len returns the number of registered route.
 func (txn *Txn) Len() int {
+	if txn.rootTxn == nil {
+		panic(ErrSettledTxn)
+	}
 	return txn.rootTxn.size
 }