Scalar-replace nonescaping local atoms

Author: jfhamlin

benchmark/aot/fixtures/constant-arithmetic.glj

@@ -8,6 +8,13 @@
 (defn count-input [value]
   (count value))
 
+(defn local-counter []
+  (let [state (atom 0)]
+    (swap! state inc)
+    (swap! state #(+ % 40))
+    (reset! state (+ @state 1))
+    @state))
+
 (defn run []
   (loop [i 0
          total 0]

benchmark/aot/main.go

@@ -205,6 +205,10 @@ func writeBenchmark(temp string, fixtures []fixture) {
 	source.WriteString("\tif got := lang.Apply1(countCaller, lang.NewVector(1, 2, 3)); !lang.Equals(got, int64(99)) {\n")
 	source.WriteString("\t\tt.Fatalf(\"direct core call ignored count redefinition: got %v, want 99\", got)\n")
 	source.WriteString("\t}\n")
+	source.WriteString("\tlocalCounter := constant_arithmeticNS.FindInternedVar(lang.NewSymbol(\"local-counter\")).Get().(lang.IFn)\n")
+	source.WriteString("\tif got := localCounter.Invoke(); !lang.Equals(got, int64(42)) {\n")
+	source.WriteString("\t\tt.Fatalf(\"scalar-replaced local atom = %v, want 42\", got)\n")
+	source.WriteString("\t}\n")
 	source.WriteString("}\n\n")
 	for _, fixture := range fixtures {
 		benchmarkName := exportedName(strings.TrimPrefix(fixture.nsName, "bench."))

pkg/lang/atom.go

@@ -119,6 +119,10 @@ func (a *Atom) CompareAndSet(oldv, newv interface{}) bool {
 
 func (a *Atom) compareAndSetBox(old *Box, newv interface{}) bool {
 	// TODO: validate
+	if Identical(old.val, newv) &&
+		(a.watches == nil || a.watches.Count() == 0) {
+		return a.state.CompareAndSwap(old, old)
+	}
 	swapped := a.state.CompareAndSwap(old, NewBox(newv))
 	if swapped {
 		a.notifyWatches(old.val, newv)
@@ -128,10 +132,19 @@ func (a *Atom) compareAndSetBox(old *Box, newv interface{}) bool {
 
 func (a *Atom) Reset(newVal interface{}) interface{} {
 	// TODO: validate
-
-	old := a.state.Swap(NewBox(newVal))
-	a.notifyWatches(old.val, newVal)
-	return newVal
+	for {
+		old := a.state.Load()
+		if Identical(old.val, newVal) &&
+			(a.watches == nil || a.watches.Count() == 0) {
+			if a.state.CompareAndSwap(old, old) {
+				return newVal
+			}
+			continue
+		}
+		old = a.state.Swap(NewBox(newVal))
+		a.notifyWatches(old.val, newVal)
+		return newVal
+	}
 }
 
 func (a *Atom) Meta() IPersistentMap {

pkg/lang/atom_test.go

@@ -29,3 +29,42 @@ func TestAtomFixedAritySwapAllocatesOnlyNewState(t *testing.T) {
 		t.Fatalf("Swap0 allocated %v objects, want at most 1", got)
 	}
 }
+
+func TestAtomIdenticalUpdatesAvoidNewStateAllocation(t *testing.T) {
+	value := &struct{}{}
+	atom := NewAtom(value)
+
+	if got := testing.AllocsPerRun(1_000, func() {
+		atom.Reset(value)
+	}); got != 0 {
+		t.Fatalf("identical Reset allocated %v objects, want 0", got)
+	}
+	if got := testing.AllocsPerRun(1_000, func() {
+		if !atom.CompareAndSet(value, value) {
+			t.Fatal("identical CompareAndSet failed")
+		}
+	}); got != 0 {
+		t.Fatalf("identical CompareAndSet allocated %v objects, want 0", got)
+	}
+}
+
+func TestAtomIdenticalUpdatesStillNotifyWatches(t *testing.T) {
+	value := &struct{}{}
+	atom := NewAtom(value)
+	calls := 0
+	atom.AddWatch("watch", FnFunc4(func(_, _, oldValue, newValue any) any {
+		calls++
+		if oldValue != value || newValue != value {
+			t.Fatal("watch received the wrong value")
+		}
+		return nil
+	}))
+
+	atom.Reset(value)
+	if !atom.CompareAndSet(value, value) {
+		t.Fatal("identical CompareAndSet failed")
+	}
+	if calls != 2 {
+		t.Fatalf("watch called %d times, want 2", calls)
+	}
+}

pkg/runtime/codegen.go

@@ -32,8 +32,9 @@ import (
 
 // varScope represents a variable allocation scope
 type varScope struct {
-	nextNum int
-	names   map[string]string // maps Clojure names to Go variable names
+	nextNum    int
+	names      map[string]string // maps Clojure names to Go variable names
+	localAtoms map[string]bool   // local atoms proven not to escape
 }
 
 // recurContext represents the context for a loop/recur form
@@ -1669,6 +1670,9 @@ func (g *Generator) generateVarDeref(node *ast.Node) string {
 // generateInvoke generates code for an Invoke node
 func (g *Generator) generateInvoke(node *ast.Node) string {
 	invokeNode := node.Sub.(*ast.InvokeNode)
+	if result, ok := g.generateLocalAtomInvoke(invokeNode); ok {
+		return result
+	}
 	if plan := analyzeReducePipeline(invokeNode); plan != nil {
 		return g.generateAOTReducePipeline(invokeNode, plan)
 	}
@@ -2031,7 +2035,7 @@ func (g *Generator) generateLet(node *ast.Node, isLoop bool) string {
 	}
 
 	// Emit bindings directly to g.w
-	for _, binding := range letNode.Bindings {
+	for bindingIndex, binding := range letNode.Bindings {
 		bindingNode := binding.Sub.(*ast.BindingNode)
 		name := bindingNode.Name.Name()
 		init := bindingNode.Init
@@ -2040,9 +2044,25 @@ func (g *Generator) generateLet(node *ast.Node, isLoop bool) string {
 		g.writef("// let binding \"%s\"\n", name)
 
 		// Generate initialization code
-		initCode := g.generateASTNode(init)
+		var localAtomInit *ast.Node
+		if !isLoop {
+			localAtomInit = scalarReplaceableAtomInit(
+				bindingNode,
+				letNode.Bindings[bindingIndex+1:],
+				letNode.Body,
+			)
+		}
+		initCode := ""
+		if localAtomInit != nil {
+			initCode = g.generateASTNode(localAtomInit)
+		} else {
+			initCode = g.generateASTNode(init)
+		}
 		varName := g.allocateLocal(name)
 		g.writef("var %s any = %s\n", varName, initCode)
+		if localAtomInit != nil {
+			g.markLocalAtom(name)
+		}
 		g.writeAssign("_", varName) // Prevent unused variable warning
 
 		// Collect binding variables for loop
@@ -2754,8 +2774,9 @@ func (g *Generator) pushVarScope() {
 
 	// Push new scope onto the stack
 	g.varScopes = append(g.varScopes, varScope{
-		nextNum: nextNum,
-		names:   make(map[string]string),
+		nextNum:    nextNum,
+		names:      make(map[string]string),
+		localAtoms: make(map[string]bool),
 	})
 }
 
@@ -2866,6 +2887,20 @@ func (g *Generator) getLocal(name string) string {
 	panic(fmt.Sprintf("variable %s not found in any scope", name))
 }
 
+func (g *Generator) markLocalAtom(name string) {
+	g.varScopes[len(g.varScopes)-1].localAtoms[name] = true
+}
+
+func (g *Generator) getLocalAtom(name string) (string, bool) {
+	for i := len(g.varScopes) - 1; i >= 0; i-- {
+		scope := &g.varScopes[i]
+		if varName, ok := scope.names[name]; ok {
+			return varName, scope.localAtoms[name]
+		}
+	}
+	return "", false
+}
+
 // allocateTempVar allocates a fresh temporary variable without name tracking
 func (g *Generator) allocateTempVar() string {
 	if len(g.varScopes) == 0 {

pkg/runtime/codegen_direct.go

@@ -204,7 +204,8 @@ func (g *Generator) aotExternalInvokeTarget(
 	}
 	arity := len(invoke.Args)
 	intrinsic := aotExternalIntrinsic(vr, arity)
-	if arity > 5 || !aotSupportsArity(codegenVarValue(vr), arity) {
+	if arity > 5 ||
+		(intrinsic == "" && !aotSupportsArity(codegenVarValue(vr), arity)) {
 		return nil
 	}
 	key := aotExternalCallKey{vr: vr, arity: arity}
@@ -236,11 +237,13 @@ func aotExternalIntrinsic(vr *lang.Var, arity int) string {
 	switch {
 	case name == "assoc" && (arity == 3 || arity == 5):
 	case name == "count" && arity == 1:
+	case name == "dec" && arity == 1:
 	case name == "cons" && arity == 2:
 	case name == "conj" && arity == 2:
 	case name == "empty?" && arity == 1:
 	case name == "first" && arity == 1:
 	case name == "get" && (arity == 2 || arity == 3):
+	case name == "inc" && arity == 1:
 	case name == "next" && arity == 1:
 	case name == "nth" && (arity == 2 || arity == 3):
 	case name == "peek" && arity == 1:
@@ -267,6 +270,8 @@ func (g *Generator) aotExternalIntrinsicCall(
 		)
 	case "count":
 		return fmt.Sprintf("lang.Count(%s)", args[0])
+	case "dec":
+		return fmt.Sprintf("lang.Numbers.Dec(%s)", args[0])
 	case "cons":
 		return fmt.Sprintf("lang.NewCons(%s, %s)", args[0], args[1])
 	case "conj":
@@ -280,6 +285,8 @@ func (g *Generator) aotExternalIntrinsicCall(
 			return fmt.Sprintf("lang.Get(%s, %s)", args[0], args[1])
 		}
 		return fmt.Sprintf("lang.GetDefault(%s, %s, %s)", args[0], args[1], args[2])
+	case "inc":
+		return fmt.Sprintf("lang.Numbers.Inc(%s)", args[0])
 	case "next":
 		return fmt.Sprintf("lang.Next(%s)", args[0])
 	case "nth":