From 634c0b5175c7a85216a59a94d52a3765aa6aa0fa Mon Sep 17 00:00:00 2001
From: Teemu Koponen <koponen@styra.com>
Date: Wed, 1 Nov 2023 15:30:43 -0700
Subject: [PATCH] internal/planner: Insert general ref head objects starting
 from the leaves, not root.

This way the object insert operations can return a new object instance.

Before, the object construction for a rule like

    p[a][b] := ...

would look like this:

    *ir.BlockStmt BlockStmt (1 blocks)
      *ir.Block Block (3 statements)
        *ir.BlockStmt BlockStmt (1 blocks)
          *ir.Block Block (2 statements)
            *ir.DotStmt &{Source:{Value:Local<2>} Key:{Value:Local<10>} Target:Local<14>}
            *ir.BreakStmt &{Index:1}
        *ir.MakeObjectStmt &{Target:Local<14>}
        *ir.ObjectInsertOnceStmt &{Key:{Value:Local<10>} Value:{Value:Local<14>} Object:Local<2>}
    *ir.ObjectInsertOnceStmt &{Key:{Value:Local<11>} Value:{Value:Local<13>} Object:Local<14>}

Now, it'll look like

    *ir.BlockStmt BlockStmt (1 blocks)
      *ir.Block Block (2 statements)
        *ir.BlockStmt BlockStmt (1 blocks)
          *ir.Block Block (2 statements)
            *ir.DotStmt &{Source:{Value:Local<2>} Key:{Value:Local<10>} Target:Local<14>}
            *ir.BreakStmt &{Index:1}
        *ir.MakeObjectStmt &{Target:Local<14>}
    *ir.ObjectInsertOnceStmt &{Key:{Value:Local<11>} Value:{Value:Local<13>} Object:Local<14>}
    *ir.ObjectInsertStmt &{Key:{Value:Local<10>} Value:{Value:Local<14>} Object:Local<2>}

so the object in Local<14> is built first, and the added to object Local<2>.

Co-authored-by: Teemu Koponen <koponen@styra.com>
Signed-off-by: Stephan Renatus <stephan@styra.com>
---
 internal/planner/planner.go                   | 10 +++--
 .../refheads/test-refs-as-rule-heads.yaml     | 38 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/internal/planner/planner.go b/internal/planner/planner.go
index 87b64e8671..a6405291e7 100644
--- a/internal/planner/planner.go
+++ b/internal/planner/planner.go
@@ -464,7 +464,8 @@ func (p *Planner) planDotOr(obj ir.Local, key ir.Operand, or stmtFactory, iter p
 	// | | | dot &{Source:Local<obj> Key:{Value:Local<key>} Target:Local<val>}
 	// | | | break 1
 	// | | or &{Target:Local<val>}
-	// | | *ir.ObjectInsertOnceStmt &{Key:{Value:Local<key>} Value:{Value:Local<val>} Object:Local<obj>}
+	// | iter &{Target:Local<val>} # may update Local<val>.
+	// | *ir.ObjectInsertStmt &{Key:{Value:Local<key>} Value:{Value:Local<val>} Object:Local<obj>}
 
 	prev := p.curr
 	dotBlock := &ir.Block{}
@@ -482,13 +483,16 @@ func (p *Planner) planDotOr(obj ir.Local, key ir.Operand, or stmtFactory, iter p
 		Stmts: []ir.Stmt{
 			&ir.BlockStmt{Blocks: []*ir.Block{dotBlock}}, // FIXME: Set Location
 			or(val),
-			&ir.ObjectInsertOnceStmt{Key: key, Value: op(val), Object: obj},
 		},
 	}
 
 	p.curr = prev
 	p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{outerBlock}})
-	return iter(val)
+	if err := iter(val); err != nil {
+		return err
+	}
+	p.appendStmt(&ir.ObjectInsertStmt{Key: key, Value: op(val), Object: obj})
+	return nil
 }
 
 func (p *Planner) planNestedObjects(obj ir.Local, ref ast.Ref, iter planLocalIter) error {
diff --git a/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml b/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml
index ebe69d0027..e289133336 100644
--- a/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml
+++ b/test/cases/testdata/refheads/test-refs-as-rule-heads.yaml
@@ -320,4 +320,40 @@ cases:
       public: false
   want_result:
   - x:
-    - n1
\ No newline at end of file
+    - n1
+- modules:
+  - |
+    package test
+    import future.keywords
+
+    p[a][b][c][d][e] if {
+        some a in numbers.range(1, 5)
+        some b in numbers.range(1, 5)
+        some c in numbers.range(1, 5)
+        some d in numbers.range(1, 5)
+        some e in numbers.range(1, 5)
+        a+b+c+d+e == 24
+    }
+  note: refheads/many-vars
+  query: data.test.p = x
+  want_result:
+  - x:
+      4:
+        5:
+          5:
+            5:
+              5: true
+      5:
+        4:
+          5:
+            5:
+              5: true
+        5:
+          4:
+            5:
+              5: true
+          5:
+            4:
+              5: true
+            5:
+              4: true