layout: avoid heap escapes in Stack and Flex

Stack.Layout and Flex.Layout caused a lot of heap allocations / escapes. The
reason was that scratch space for dims and call and was inside Stack/FlexChild.
child.call.Add(gtx.Ops) confused the go escape analysis and caused the entired
children slice to escape to the heap, including all widgets in it. This caused a
lot of heap allocations. Now the scratch space is separate from children, and
for cases len(children) <= 32, we will allocate the scratch space on the stack.
For cases len(children) > 32, only the scratch space gets allocated from the
heap, during append.

Signed-off-by: vsariola <5684185+vsariola@users.noreply.github.com>

layout/flex.go

@@ -30,10 +30,6 @@ type FlexChild struct {
 	weight float32
 
 	widget Widget
-
-	// Scratch space.
-	call op.CallOp
-	dims Dimensions
 }
 
 // Spacing determine the spacing mode for a Flex.
@@ -88,6 +84,20 @@ func (f Flex) Layout(gtx Context, children ...FlexChild) Dimensions {
 	remaining := mainMax
 	var totalWeight float32
 	cgtx := gtx
+	// Note: previously the scratch space was inside FlexChild.
+	// child.call.Add(gtx.Ops) confused the go escape analysis and caused the
+	// entired children slice to be allocated on the heap, including all widgets
+	// in it. This produced a lot of object allocations. Now the scratch space
+	// is separate from children, and for cases len(children) <= 32, we will
+	// allocate the scratch space on the stack. For cases len(children) > 32,
+	// only the scratch space gets allocated from the heap, during append.
+	type scratchSpace struct {
+		call op.CallOp
+		dims Dimensions
+	}
+	var scratchArray [32]scratchSpace
+	scratch := scratchArray[:0]
+	scratch = append(scratch, make([]scratchSpace, len(children))...)
 	// Lay out Rigid children.
 	for i, child := range children {
 		if child.flex {
@@ -104,8 +114,8 @@ func (f Flex) Layout(gtx Context, children ...FlexChild) Dimensions {
 		if remaining < 0 {
 			remaining = 0
 		}
-		children[i].call = c
-		children[i].dims = dims
+		scratch[i].call = c
+		scratch[i].dims = dims
 	}
 	if w := f.WeightSum; w != 0 {
 		totalWeight = w
@@ -139,16 +149,16 @@ func (f Flex) Layout(gtx Context, children ...FlexChild) Dimensions {
 		if remaining < 0 {
 			remaining = 0
 		}
-		children[i].call = c
-		children[i].dims = dims
+		scratch[i].call = c
+		scratch[i].dims = dims
 	}
 	maxCross := crossMin
 	var maxBaseline int
-	for _, child := range children {
-		if c := f.Axis.Convert(child.dims.Size).Y; c > maxCross {
+	for _, scratchChild := range scratch {
+		if c := f.Axis.Convert(scratchChild.dims.Size).Y; c > maxCross {
 			maxCross = c
 		}
-		if b := child.dims.Size.Y - child.dims.Baseline; b > maxBaseline {
+		if b := scratchChild.dims.Size.Y - scratchChild.dims.Baseline; b > maxBaseline {
 			maxBaseline = b
 		}
 	}
@@ -169,8 +179,8 @@ func (f Flex) Layout(gtx Context, children ...FlexChild) Dimensions {
 			mainSize += space / (len(children) * 2)
 		}
 	}
-	for i, child := range children {
-		dims := child.dims
+	for i, scratchChild := range scratch {
+		dims := scratchChild.dims
 		b := dims.Size.Y - dims.Baseline
 		var cross int
 		switch f.Alignment {
@@ -185,7 +195,7 @@ func (f Flex) Layout(gtx Context, children ...FlexChild) Dimensions {
 		}
 		pt := f.Axis.Convert(image.Pt(mainSize, cross))
 		trans := op.Offset(pt).Push(gtx.Ops)
-		child.call.Add(gtx.Ops)
+		scratchChild.call.Add(gtx.Ops)
 		trans.Pop()
 		mainSize += f.Axis.Convert(dims.Size).X
 		if i < len(children)-1 {

layout/stack.go

@@ -20,10 +20,6 @@ type Stack struct {
 type StackChild struct {
 	expanded bool
 	widget   Widget
-
-	// Scratch space.
-	call op.CallOp
-	dims Dimensions
 }
 
 // Stacked returns a Stack child that is laid out with no minimum
@@ -52,6 +48,20 @@ func (s Stack) Layout(gtx Context, children ...StackChild) Dimensions {
 	// First lay out Stacked children.
 	cgtx := gtx
 	cgtx.Constraints.Min = image.Point{}
+	// Note: previously the scratch space was inside StackChild.
+	// child.call.Add(gtx.Ops) confused the go escape analysis and caused the
+	// entired children slice to be allocated on the heap, including all widgets
+	// in it. This produced a lot of object allocations. Now the scratch space
+	// is separate from children, and for cases len(children) <= 32, we will
+	// allocate the scratch space on the stack. For cases len(children) > 32,
+	// only the scratch space gets allocated from the heap, during append.
+	type scratchSpace struct {
+		call op.CallOp
+		dims Dimensions
+	}
+	var scratchArray [32]scratchSpace
+	scratch := scratchArray[:0]
+	scratch = append(scratch, make([]scratchSpace, len(children))...)
 	for i, w := range children {
 		if w.expanded {
 			continue
@@ -65,8 +75,8 @@ func (s Stack) Layout(gtx Context, children ...StackChild) Dimensions {
 		if h := dims.Size.Y; h > maxSZ.Y {
 			maxSZ.Y = h
 		}
-		children[i].call = call
-		children[i].dims = dims
+		scratch[i].call = call
+		scratch[i].dims = dims
 	}
 	// Then lay out Expanded children.
 	for i, w := range children {
@@ -83,14 +93,14 @@ func (s Stack) Layout(gtx Context, children ...StackChild) Dimensions {
 		if h := dims.Size.Y; h > maxSZ.Y {
 			maxSZ.Y = h
 		}
-		children[i].call = call
-		children[i].dims = dims
+		scratch[i].call = call
+		scratch[i].dims = dims
 	}
 
 	maxSZ = gtx.Constraints.Constrain(maxSZ)
 	var baseline int
-	for _, ch := range children {
-		sz := ch.dims.Size
+	for _, scratchChild := range scratch {
+		sz := scratchChild.dims.Size
 		var p image.Point
 		switch s.Alignment {
 		case N, S, Center:
@@ -105,10 +115,10 @@ func (s Stack) Layout(gtx Context, children ...StackChild) Dimensions {
 			p.Y = maxSZ.Y - sz.Y
 		}
 		trans := op.Offset(p).Push(gtx.Ops)
-		ch.call.Add(gtx.Ops)
+		scratchChild.call.Add(gtx.Ops)
 		trans.Pop()
 		if baseline == 0 {
-			if b := ch.dims.Baseline; b != 0 {
+			if b := scratchChild.dims.Baseline; b != 0 {
 				baseline = b + maxSZ.Y - sz.Y - p.Y
 			}
 		}