From bcb722925dca1b8c190b9bd1dbb9663d45adfb65 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Fri, 12 Mar 2021 14:30:32 +0100
Subject: [PATCH] wasm: fix offset bug

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 internal/compiler/wasm/optimizations.go      | 16 ++----
 internal/compiler/wasm/optimizations_test.go | 56 ++++++++++++++++++++
 2 files changed, 59 insertions(+), 13 deletions(-)
 create mode 100644 internal/compiler/wasm/optimizations_test.go

diff --git a/internal/compiler/wasm/optimizations.go b/internal/compiler/wasm/optimizations.go
index 95c33e8f83..3873cce186 100644
--- a/internal/compiler/wasm/optimizations.go
+++ b/internal/compiler/wasm/optimizations.go
@@ -193,16 +193,6 @@ func (c *Compiler) removeUnusedCode() error {
 	}
 	c.module.Names.Functions = funcNames
 
-	// functions we've compiled only get a new index
-	funcs := []funcCode{}
-	for _, f := range c.funcsCode {
-		oldIdx := c.funcs[f.name]
-		if _, ok := keepFuncs[oldIdx]; ok {
-			funcs = append(funcs, f)
-		}
-	}
-	c.funcsCode = funcs
-
 	// For anything that we don't want, replace the function code entries'
 	// expressions with `unreachable`.
 	// We do this because it lets the resulting wasm module pass `wasm-validate`,
@@ -217,9 +207,9 @@ func (c *Compiler) removeUnusedCode() error {
 		return fmt.Errorf("write code entry: %w", err)
 	}
 	for i := range c.module.Code.Segments {
-		if _, ok := keepFuncs[uint32(i)]; !ok {
-			idx := i - c.functionImportCount()
-			c.module.Code.Segments[idx].Code = buf.Bytes()
+		idx := i + c.functionImportCount()
+		if _, ok := keepFuncs[uint32(idx)]; !ok {
+			c.module.Code.Segments[i].Code = buf.Bytes()
 		}
 	}
 	return nil
diff --git a/internal/compiler/wasm/optimizations_test.go b/internal/compiler/wasm/optimizations_test.go
new file mode 100644
index 0000000000..551b770c21
--- /dev/null
+++ b/internal/compiler/wasm/optimizations_test.go
@@ -0,0 +1,56 @@
+package wasm
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/internal/planner"
+)
+
+func TestRemoveUnusedCode(t *testing.T) {
+	policy, err := planner.New().
+		WithQueries([]planner.QuerySet{
+			{
+				Name:    "test",
+				Queries: []ast.Body{ast.MustParseBody(`input.foo = 1`)},
+			},
+		}).Plan()
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	c := New().WithPolicy(policy)
+	mod, err := c.Compile()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// NOTE(sr): our unused code elimination has the following invariant:
+	// if a function is not used, both its code and its name are removed
+	// from the code section, and name sections respectively.
+	idxToName := map[uint32]string{}
+	for _, nm := range mod.Names.Functions {
+		idxToName[nm.Index] = nm.Name
+	}
+	for i, seg := range mod.Code.Segments {
+		idx := i + c.functionImportCount()
+		noop := len(seg.Code) == 3
+		name, ok := idxToName[uint32(idx)]
+		if noop && ok {
+			t.Errorf("func[%d] has name (%s) and no code", idx, name)
+		}
+		if !noop && !ok {
+			t.Errorf("func[%d] has code but no name", idx)
+		}
+	}
+
+	// Having established that, we can check that this simple policy
+	// has no re2-related code by consulting the name map:
+	for _, nm := range mod.Names.Functions {
+		if strings.Contains(nm.Name, "re2") {
+			t.Errorf("expected no re2-related functions, found %s", nm.Name)
+		}
+	}
+}