wasm: plan data lookup for ref when call_indirect mapping lookup fails (#3344)

If there is no data function for the vars (known at runtime), we'll have
to consult `data`. Before, this was ignored.

Fixes #3305.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

internal/compiler/wasm/wasm.go

@@ -1375,12 +1375,12 @@ func (c *Compiler) compileUpsert(local ir.Local, path []int, value ir.LocalOrCon
 }
 
 func (c *Compiler) compileCallDynamicStmt(stmt *ir.CallDynamicStmt, result *[]instruction.Instruction) error {
-	block := instruction.Block{}
+	instrs := []instruction.Instruction{}
 	larray := c.genLocal()
 	lidx := c.genLocal()
 
 	// init array:
-	block.Instrs = append(block.Instrs,
+	instrs = append(instrs,
 		instruction.I32Const{Value: int32(len(stmt.Path))},
 		instruction.Call{Index: c.function(opaArrayWithCap)},
 		instruction.SetLocal{Index: larray},
@@ -1388,7 +1388,7 @@ func (c *Compiler) compileCallDynamicStmt(stmt *ir.CallDynamicStmt, result *[]in
 
 	// append to it:
 	for _, lv := range stmt.Path {
-		block.Instrs = append(block.Instrs,
+		instrs = append(instrs,
 			instruction.GetLocal{Index: larray},
 			c.instrRead(lv),
 			instruction.Call{Index: c.function(opaArrayAppend)},
@@ -1397,7 +1397,7 @@ func (c *Compiler) compileCallDynamicStmt(stmt *ir.CallDynamicStmt, result *[]in
 
 	// prep stack for later call_indirect
 	for _, arg := range stmt.Args {
-		block.Instrs = append(block.Instrs, instruction.GetLocal{Index: c.local(arg)})
+		instrs = append(instrs, instruction.GetLocal{Index: c.local(arg)})
 	}
 
 	tpe := module.FunctionType{
@@ -1406,22 +1406,22 @@ func (c *Compiler) compileCallDynamicStmt(stmt *ir.CallDynamicStmt, result *[]in
 	}
 	typeIndex := c.emitFunctionType(tpe)
 
-	block.Instrs = append(block.Instrs,
+	instrs = append(instrs,
 		// lookup elem idx via larray path
 		instruction.GetLocal{Index: larray},
 		instruction.Call{Index: c.function(opaMappingLookup)}, // [arg0 arg1 larray] -> [arg0 arg1 tbl_idx]
 		instruction.TeeLocal{Index: lidx},
-		instruction.I32Eqz{}, // mapping not found
-		instruction.BrIf{Index: 1},
+		instruction.I32Eqz{},       // mapping not found
+		instruction.BrIf{Index: 0}, // check data
 
 		instruction.GetLocal{Index: lidx},
 		instruction.CallIndirect{Index: typeIndex}, // [arg0 arg1 tbl_idx] -> [res]
 		instruction.TeeLocal{Index: c.local(stmt.Result)},
 		instruction.I32Eqz{},
-		instruction.BrIf{Index: 1},
+		instruction.BrIf{Index: 2}, // mapping found, "undefined" result counts
 	)
 
-	*result = append(*result, block)
+	*result = append(*result, instrs...)
 	return nil
 }
 

internal/planner/planner.go

@@ -1463,16 +1463,77 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
 				}
 			}
 
-			p.ltarget = p.newLocal()
+			// We're planning a structure like this:
+			//
+			// block a
+			//   block b
+			//     block c1
+			//       opa_mapping_lookup || br c1
+			//       call_indirect      || br a
+			//       br b
+			//     end
+			//     block c2
+			//       dot i   || br c2
+			//       dot i+1 || br c2
+			//       br b
+			//     end
+			//     br a
+			//   end
+			//   dot i+2 || br a
+			//   dot i+3 || br a
+			// end
+			//
+			// We have to do it like this because the dot IR stmts
+			// are compiled to `br 0`, the innermost block, if they
+			// fail.
+			// The "c2" block will construct the reference from `data`
+			// only, in case the mapping lookup doesn't yield a func
+			// to call_dynamic.
+
+			ltarget := p.newLocal()
+			p.ltarget = ltarget
+			prev := p.curr
+
+			callDynBlock := &ir.Block{} // "c1" in the sketch
+			p.curr = callDynBlock
 			p.appendStmt(&ir.CallDynamicStmt{
 				Args: []ir.Local{
 					p.vars.GetOrEmpty(ast.InputRootDocument.Value.(ast.Var)),
 					p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var)),
 				},
 				Path:   path,
-				Result: p.ltarget.(ir.Local),
+				Result: ltarget,
+			})
+			p.appendStmt(&ir.BreakStmt{Index: 1})
+
+			dotBlock := &ir.Block{} // "c2" in the sketch above
+			p.curr = dotBlock
+			p.ltarget = p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var))
+
+			return p.planRefRec(ref[:index+1], 1, func() error {
+				p.appendStmt(&ir.AssignVarStmt{
+					Source: p.ltarget.(ir.Local),
+					Target: ltarget,
+				})
+				p.appendStmt(&ir.BreakStmt{Index: 1})
+				p.ltarget = ltarget
+
+				outerBlock := &ir.Block{Stmts: []ir.Stmt{
+					&ir.BlockStmt{Blocks: []*ir.Block{
+						{ // block "b" in the sketch above
+							Stmts: []ir.Stmt{
+								&ir.BlockStmt{Blocks: []*ir.Block{callDynBlock, dotBlock}},
+								&ir.BreakStmt{Index: 1}},
+						},
+					}},
+				}}
+				p.curr = outerBlock
+				return p.planRefRec(ref, index+1, func() error { // rest of the ref
+					p.curr = prev
+					p.appendStmt(&ir.BlockStmt{Blocks: []*ir.Block{outerBlock}})
+					return iter()
+				})
 			})
-			return p.planRefRec(ref, index+1, iter)
 		}
 	}
 

test/cases/testdata/negation/test-negation-data-ref-with-var.yaml

@@ -51,6 +51,22 @@ cases:
   - |
     package bar
     p = 7
-  note: 'negation/neg: ref with variable and virtual doc'
+  note: 'negation/neg: ref with variable (hit) and virtual doc (miss))'
+  query: data.foo.p = x
+  want_result: []
+- data:
+    bar:
+      q: 8
+  modules:
+  - |
+    package foo
+    p {
+      k := "q"
+      not data.bar[k]
+    }
+  - |
+    package bar
+    p = 7
+  note: 'negation/neg: ref with variable (miss) and virtual doc (hit)'
   query: data.foo.p = x
   want_result: []

test/wasm/assets/019_call_indirect_optimization.yaml

@@ -85,3 +85,39 @@ cases:
       - x: q
         z: 1
         w: 1
+  - note: data ref used after failed lookup
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo
+        p = y {
+          k := "q"
+          y := data.bar[k].baz
+        }
+      - |
+        package bar
+        p = { "baz": 100 }
+    data:
+      bar:
+        q:
+          baz: 8
+    want_result:
+      - x: 8
+  - note: data ref not used after successful lookup
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo
+        p = y {
+          k := "p"
+          y := data.bar[k].baz
+        }
+      - |
+        package bar
+        p = { "baz": 100 }
+    data:
+      bar:
+        q:
+          baz: 8
+    want_result:
+      - x: 100