From 09a9af8d2beb66c93e8b9b900430c7ccadb0c734 Mon Sep 17 00:00:00 2001
From: Adrian Hesketh <adrianhesketh@hushmail.com>
Date: Wed, 1 Jan 2025 16:17:16 +0000
Subject: [PATCH] feat: support concurrent children rendering

---
 children.go       | 37 ++++++++++++++++++++
 once.go           |  5 ++-
 runtime.go        | 87 +++++++++++++++++------------------------------
 scripttemplate.go |  3 +-
 4 files changed, 72 insertions(+), 60 deletions(-)
 create mode 100644 children.go

diff --git a/children.go b/children.go
new file mode 100644
index 000000000..b2882cc6f
--- /dev/null
+++ b/children.go
@@ -0,0 +1,37 @@
+package templ
+
+import "context"
+
+type contextChildrenType int
+
+const contextChildrenKey = contextChildrenType(1)
+
+type contextChildrenValue struct {
+	children Component
+}
+
+func WithChildren(ctx context.Context, children Component) context.Context {
+	ctx = InitializeContext(ctx)
+	if children == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, contextChildrenKey, &contextChildrenValue{children: children})
+}
+
+func GetChildren(ctx context.Context) Component {
+	if ctx == nil {
+		return NopComponent
+	}
+	v, ok := ctx.Value(contextChildrenKey).(*contextChildrenValue)
+	if !ok || v.children == nil {
+		return NopComponent
+	}
+	return v.children
+}
+
+func ClearChildren(ctx context.Context) context.Context {
+	if ctx == nil || ctx.Value(contextChildrenKey) == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, contextChildrenKey, nil)
+}
diff --git a/once.go b/once.go
index 7860ab840..6c721fea5 100644
--- a/once.go
+++ b/once.go
@@ -51,11 +51,10 @@ type OnceHandle struct {
 // Once returns a component that renders its children once per context.
 func (o *OnceHandle) Once() Component {
 	return ComponentFunc(func(ctx context.Context, w io.Writer) (err error) {
-		_, v := getContext(ctx)
-		if v.getHasBeenRendered(o) {
+		ctx, v := getContext(ctx)
+		if !v.shouldRenderOnce(o) {
 			return nil
 		}
-		v.setHasBeenRendered(o)
 		if o.c != nil {
 			return o.c.Render(ctx, w)
 		}
diff --git a/runtime.go b/runtime.go
index b55459b63..6ea4a385c 100644
--- a/runtime.go
+++ b/runtime.go
@@ -53,30 +53,9 @@ func GetNonce(ctx context.Context) (nonce string) {
 	return v.nonce
 }
 
-func WithChildren(ctx context.Context, children Component) context.Context {
-	ctx, v := getContext(ctx)
-	v.children = &children
-	return ctx
-}
-
-func ClearChildren(ctx context.Context) context.Context {
-	_, v := getContext(ctx)
-	v.children = nil
-	return ctx
-}
-
 // NopComponent is a component that doesn't render anything.
 var NopComponent = ComponentFunc(func(ctx context.Context, w io.Writer) error { return nil })
 
-// GetChildren from the context.
-func GetChildren(ctx context.Context) Component {
-	_, v := getContext(ctx)
-	if v.children == nil {
-		return NopComponent
-	}
-	return *v.children
-}
-
 // EscapeString escapes HTML text within templates.
 func EscapeString(s string) string {
 	return html.EscapeString(s)
@@ -286,7 +265,7 @@ func (cssm CSSMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// Add registered classes to the context.
 	ctx, v := getContext(r.Context())
 	for _, c := range cssm.CSSHandler.Classes {
-		v.addClass(c.ID)
+		v.shouldRenderClass(c.ID)
 	}
 	// Serve the request. Templ components will use the updated context
 	// to know to skip rendering <style> elements for any component CSS
@@ -353,9 +332,8 @@ func renderCSSItemsToBuilder(sb *strings.Builder, v *contextValue, classes ...an
 	for _, c := range classes {
 		switch ccc := c.(type) {
 		case ComponentCSSClass:
-			if !v.hasClassBeenRendered(ccc.ID) {
+			if v.shouldRenderClass(ccc.ID) {
 				sb.WriteString(string(ccc.Class))
-				v.addClass(ccc.ID)
 			}
 		case KeyValue[ComponentCSSClass, bool]:
 			if !ccc.Value {
@@ -493,55 +471,52 @@ type contextKeyType int
 const contextKey = contextKeyType(0)
 
 type contextValue struct {
+	m           *sync.Mutex
 	ss          map[string]struct{}
 	onceHandles map[*OnceHandle]struct{}
-	children    *Component
 	nonce       string
 }
 
-func (v *contextValue) setHasBeenRendered(h *OnceHandle) {
+func (v *contextValue) shouldRenderOnce(h *OnceHandle) (render bool) {
+	v.m.Lock()
+	defer v.m.Unlock()
 	if v.onceHandles == nil {
-		v.onceHandles = map[*OnceHandle]struct{}{}
+		v.onceHandles = make(map[*OnceHandle]struct{})
 	}
-	v.onceHandles[h] = struct{}{}
-}
-
-func (v *contextValue) getHasBeenRendered(h *OnceHandle) (ok bool) {
-	if v.onceHandles == nil {
-		v.onceHandles = map[*OnceHandle]struct{}{}
+	_, rendered := v.onceHandles[h]
+	if rendered {
+		return false
 	}
-	_, ok = v.onceHandles[h]
-	return
+	v.onceHandles[h] = struct{}{}
+	return true
 }
 
-func (v *contextValue) addScript(s string) {
+func (v *contextValue) shouldRenderScript(s string) (render bool) {
+	v.m.Lock()
+	defer v.m.Unlock()
 	if v.ss == nil {
-		v.ss = map[string]struct{}{}
+		v.ss = make(map[string]struct{})
+	}
+	_, rendered := v.ss["script_"+s]
+	if rendered {
+		return false
 	}
 	v.ss["script_"+s] = struct{}{}
+	return true
 }
 
-func (v *contextValue) hasScriptBeenRendered(s string) (ok bool) {
+func (v *contextValue) shouldRenderClass(s string) (render bool) {
+	v.m.Lock()
+	defer v.m.Unlock()
 	if v.ss == nil {
-		v.ss = map[string]struct{}{}
+		v.ss = make(map[string]struct{})
 	}
-	_, ok = v.ss["script_"+s]
-	return
-}
-
-func (v *contextValue) addClass(s string) {
-	if v.ss == nil {
-		v.ss = map[string]struct{}{}
+	_, rendered := v.ss["class_"+s]
+	if rendered {
+		return false
 	}
 	v.ss["class_"+s] = struct{}{}
-}
-
-func (v *contextValue) hasClassBeenRendered(s string) (ok bool) {
-	if v.ss == nil {
-		v.ss = map[string]struct{}{}
-	}
-	_, ok = v.ss["class_"+s]
-	return
+	return true
 }
 
 // InitializeContext initializes context used to store internal state used during rendering.
@@ -549,7 +524,9 @@ func InitializeContext(ctx context.Context) context.Context {
 	if _, ok := ctx.Value(contextKey).(*contextValue); ok {
 		return ctx
 	}
-	v := &contextValue{}
+	v := &contextValue{
+		m: &sync.Mutex{},
+	}
 	ctx = context.WithValue(ctx, contextKey, v)
 	return ctx
 }
diff --git a/scripttemplate.go b/scripttemplate.go
index 4962cafa2..7ee97497a 100644
--- a/scripttemplate.go
+++ b/scripttemplate.go
@@ -78,9 +78,8 @@ func RenderScriptItems(ctx context.Context, w io.Writer, scripts ...ComponentScr
 	_, v := getContext(ctx)
 	sb := new(strings.Builder)
 	for _, s := range scripts {
-		if !v.hasScriptBeenRendered(s.Name) {
+		if v.shouldRenderScript(s.Name) {
 			sb.WriteString(s.Function)
-			v.addScript(s.Name)
 		}
 	}
 	if sb.Len() > 0 {