cue/ast/astutil: avoid walking to the root to reuse scopes

As suggested by Matthew, this can avoid some overhead when having
more than a few levels of scope nesting. Plus, a pointer is smaller
than a slice, so each scope is a bit smaller.

Even with the AWS schema, which tops at a scope depth of 14,
this does lead to a measurable savings in CPU cost.
On the flip side, dealing with pointers makes the code a bit trickier,
but overall it's a coin toss as we can remove two methods.

                 │    old     │               mid                │
                 │   sec/op   │   sec/op    vs base              │
    FmtAwsSchema   2.210 ± 1%   2.167 ± 1%  -1.91% (p=0.001 n=8)

                 │     old      │                mid                 │
                 │     B/op     │     B/op      vs base              │
    FmtAwsSchema   1.064Gi ± 0%   1.064Gi ± 0%  -0.00% (p=0.005 n=8)

                 │     old     │             mid              │
                 │  allocs/op  │  allocs/op   vs base         │
    FmtAwsSchema   12.92M ± 0%   12.92M ± 0%  ~ (p=0.063 n=8)

Signed-off-by: Daniel Martí <[email protected]>
Change-Id: Iea3af538bc94da349456f9e9918052a0e6ae192b
Reviewed-on: https://cue.gerrithub.io/c/cue-lang/cue/+/1225195
Reviewed-by: Matthew Sackman <[email protected]>
Unity-Result: CUE porcuepine <[email protected]>
TryBot-Result: CUEcueckoo <[email protected]>

Author: mvdan

cue/ast/astutil/resolve.go

@@ -83,15 +83,26 @@ type ErrFunc func(pos token.Pos, msg string, args ...interface{})
 // Unresolved identifiers are recorded in [ast.File.Unresolved].
 // It will not overwrite already resolved identifiers.
 func Resolve(f *ast.File, errFn ErrFunc) {
-	visitor := &scope{errFn: errFn, identFn: resolveIdent}
+	stack := make([]*scope, 0, 8)
+	visitor := &scope{
+		errFn:      errFn,
+		identFn:    resolveIdent,
+		scopeStack: &stack,
+	}
 	ast.Walk(f, visitor.Before, nil)
 }
 
 // ResolveExpr resolves all identifiers in an expression.
 // It will not overwrite already resolved values.
 func ResolveExpr(e ast.Expr, errFn ErrFunc) {
 	f := &ast.File{}
-	visitor := &scope{file: f, errFn: errFn, identFn: resolveIdent}
+	stack := make([]*scope, 0, 8)
+	visitor := &scope{
+		file:       f,
+		errFn:      errFn,
+		identFn:    resolveIdent,
+		scopeStack: &stack,
+	}
 	ast.Walk(e, visitor.Before, nil)
 }
 
@@ -109,8 +120,9 @@ type scope struct {
 	nameFn  func(name string)
 	errFn   func(p token.Pos, msg string, args ...interface{})
 
-	// scopeStack is used to reuse scope allocations. Only set on the root scope.
-	scopeStack []*scope
+	// scopeStack is used to reuse scope allocations.
+	// The pointer is shared between the root scope and all its children.
+	scopeStack *[]*scope
 }
 
 type entry struct {
@@ -119,49 +131,34 @@ type entry struct {
 	field *ast.Field // Used for LabelAliases
 }
 
-func (s *scope) rootScope() *scope {
-	// If this ever shows up in a CPU profile because we have very deeply nested scopes,
-	// we can consider a shortcut, such as a pointer to a shared slice or struct.
-	for s.outer != nil {
-		s = s.outer
-	}
-	return s
-}
-
 func (s *scope) allocScope() *scope {
-	root := s.rootScope()
-	if n := len(root.scopeStack); n > 0 {
-		scope := root.scopeStack[n-1]
-		root.scopeStack = root.scopeStack[:n-1]
+	if n := len(*s.scopeStack); n > 0 {
+		scope := (*s.scopeStack)[n-1]
+		*s.scopeStack = (*s.scopeStack)[:n-1]
 		return scope
 	}
-	return &scope{index: make(map[string]entry, 4)}
+	return &scope{
+		index:      make(map[string]entry, 4),
+		scopeStack: s.scopeStack,
+	}
 }
 
-// putScope is not a method on scope on purpose, as we don't want to repeat
-// calls to [scope.rootScope] in [scope.freeScopesUntil].
-func putScope(root, s *scope) {
+func (s *scope) freeScope() {
 	// Ensure no pointers remain, which can hold onto memory.
-	// We only reuse the index map capacity.
-	*s = scope{index: s.index}
+	// We only reuse the index map capacity, and keep the scopeStack pointer.
+	*s = scope{index: s.index, scopeStack: s.scopeStack}
 	clear(s.index)
-	root.scopeStack = append(root.scopeStack, s)
-}
-
-func (s *scope) freeScope() {
-	root := s.rootScope()
-	putScope(root, s)
+	*s.scopeStack = append(*s.scopeStack, s)
 }
 
 // freeScopesUntil frees all scopes from s up to (but not including) 'ancestor'.
 func (s *scope) freeScopesUntil(ancestor *scope) {
-	root := s.rootScope()
 	for s != ancestor {
 		if s == nil {
 			panic("ancestor scope not found")
 		}
 		next := s.outer
-		putScope(root, s)
+		s.freeScope()
 		s = next
 	}
 }

cue/ast/astutil/sanitize.go

@@ -50,22 +50,26 @@ func Sanitize(f *ast.File) error {
 	}
 
 	// Gather all names.
+	stack := make([]*scope, 0, 8)
 	s := &scope{
-		errFn:   z.errf,
-		nameFn:  z.addName,
-		identFn: z.markUsed,
+		errFn:      z.errf,
+		nameFn:     z.addName,
+		identFn:    z.markUsed,
+		scopeStack: &stack,
 	}
 	ast.Walk(f, s.Before, nil)
 	if z.errs != nil {
 		return z.errs
 	}
 
 	// Add imports and unshadow.
+	stack = stack[:0]
 	s = &scope{
-		file:    f,
-		errFn:   z.errf,
-		identFn: z.handleIdent,
-		index:   make(map[string]entry),
+		file:       f,
+		errFn:      z.errf,
+		identFn:    z.handleIdent,
+		index:      make(map[string]entry),
+		scopeStack: &stack,
 	}
 	z.fileScope = s
 	ast.Walk(f, s.Before, nil)