wasm: replace unused functions by stub (#3206)

We're in this situation: performing dead code analysis on wasm isn't too
hard, but it requires a representation of all wasm instructions: we'd need
to be able to parse the "runtime" wasm bits, i.e., what's built using llvm
from C code. When building upon that wasm module, we process the function
bodies uninterpreted -- they are all just `[]byte` to us.

This restriction lets us get by without implementing all the wasm
instructions -- we only write what we use, and read a bare minimum to work
as outlined above.

To still be able to remove dead code, this change employs a trick: at build
time, when the aforementioned runtime wasm module is compiled, we're calling
wasm-opt on it to extract its call graph. We'll use that, together with the
functions actually planned in our wasm compiler (using the subset of
instructions that we understand), to remove all unused functions from the
name section, and replace their function bodies with `unreachable`.

We cannot really remove them, since that would require reindexing all
functions; and we cannot do that without replacing the function indices at
their call sites in the "runtime" wasm module.

Another restriction to the impact of this approach is call_indirect: We
need to keep every function that's referenced in the table -- we don't know
which function might be calling them indirectly. In a follow-up, we could
record that information and use it to further reduce the code size: we know
that if none of the regex-related builtins are used, we could also stub out
the re2-related functions.

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

Makefile

@@ -180,6 +180,7 @@ wasm-lib-build:
 ifeq ($(DOCKER_RUNNING), 1)
 	@$(MAKE) -C wasm ensure-builder build
 	cp wasm/_obj/opa.wasm internal/compiler/wasm/opa/opa.wasm
+	cp wasm/_obj/callgraph.csv internal/compiler/wasm/opa/callgraph.csv
 else
 	@echo "Docker not installed or not running. Skipping OPA-WASM library build."
 endif

build/check-working-copy.sh

@@ -3,6 +3,7 @@
 EXCEPTIONS=(
   "internal/compiler/wasm/opa/opa.go"
   "internal/compiler/wasm/opa/opa.wasm"
+  "internal/compiler/wasm/opa/callgraph.csv"
 )
 
 STATUS=$(git status --porcelain)

internal/cmd/genopawasm/main.go

@@ -28,8 +28,8 @@ func main() {
 		Use:   executable,
 		Short: executable + " <opa.wasm path>",
 		RunE: func(_ *cobra.Command, args []string) error {
-			if len(args) != 1 {
-				return fmt.Errorf("provide path of opa.wasm file")
+			if len(args) != 2 {
+				return fmt.Errorf("provide path of opa.wasm and callgraph.csv files")
 			}
 			return run(params, args)
 		},
@@ -81,6 +81,16 @@ func Bytes() ([]byte, error) {
 	return ioutil.ReadAll(gr)
 }
 
+// CallGraphCSV returns a CSV representation of the
+// OPA-WASM bytecode's call graph: 'caller,callee'
+func CallGraphCSV() ([]byte, error) {
+	cg, err := gzip.NewReader(bytes.NewBuffer(gzippedCallGraphCSV))
+	if err != nil {
+		return nil, err
+	}
+	return ioutil.ReadAll(cg)
+}
+
 `))
 
 	if err != nil {
@@ -92,7 +102,34 @@ func Bytes() ([]byte, error) {
 		return err
 	}
 
-	in, err := os.Open(args[0])
+	if err := output(out, args[0]); err != nil {
+		return err
+	}
+
+	if _, err := out.Write([]byte(`")
+`)); err != nil {
+		return err
+	}
+
+	_, err = out.Write([]byte(`var gzippedCallGraphCSV = []byte("`))
+	if err != nil {
+		return err
+	}
+
+	if err := output(out, args[1]); err != nil {
+		return err
+	}
+
+	if _, err := out.Write([]byte(`")
+`)); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func output(out io.Writer, filename string) error {
+	in, err := os.Open(filename)
 	if err != nil {
 		return err
 	}
@@ -116,14 +153,7 @@ func Bytes() ([]byte, error) {
 			return err
 		}
 	}
-
-	_, err = out.Write([]byte(`")`))
-	if err != nil {
-		return err
-	}
-
-	_, err = out.Write([]byte("\n"))
-	return err
+	return nil
 }
 
 var digits = "0123456789ABCDEF"

internal/compiler/wasm/optimizations.go

@@ -3,14 +3,18 @@ package wasm
 import (
 	"bytes"
 	"context"
+	"encoding/csv"
 	"fmt"
 	"os"
 	"os/exec"
+	"strconv"
 	"strings"
 	"time"
 
+	"github.com/open-policy-agent/opa/internal/compiler/wasm/opa"
 	"github.com/open-policy-agent/opa/internal/wasm/encoding"
 	"github.com/open-policy-agent/opa/internal/wasm/instruction"
+	"github.com/open-policy-agent/opa/internal/wasm/module"
 )
 
 const warning = `---------------------------------------------------------------
@@ -101,3 +105,140 @@ func withControlInstr(is []instruction.Instruction) bool {
 	}
 	return false
 }
+
+func unquote(s string) (string, error) {
+	return strconv.Unquote("\"" + strings.ReplaceAll(s, `\`, `\x`) + "\"")
+}
+
+func (c *Compiler) removeUnusedCode() error {
+	cgCSV, err := opa.CallGraphCSV()
+	if err != nil {
+		return fmt.Errorf("csv unpack: %w", err)
+	}
+	r := csv.NewReader(bytes.NewReader(cgCSV))
+	r.LazyQuotes = true
+	cg, err := r.ReadAll()
+	if err != nil {
+		return fmt.Errorf("csv read: %w", err)
+	}
+
+	cgIdx := map[uint32][]uint32{}
+	for i := range cg {
+		callerName, err := unquote(cg[i][0])
+		if err != nil {
+			return fmt.Errorf("unquote caller name %s: %w", cg[i][0], err)
+		}
+		calleeName, err := unquote(cg[i][1])
+		if err != nil {
+			return fmt.Errorf("unquote callee name %s: %w", cg[i][1], err)
+		}
+		caller, ok := c.funcs[callerName]
+		if !ok {
+			return fmt.Errorf("caller not found: %s (%s)", cg[i][0], callerName)
+		}
+		callee, ok := c.funcs[calleeName]
+		if !ok {
+			return fmt.Errorf("callee not found: %s (%s)", cg[i][1], calleeName)
+		}
+		cgIdx[caller] = append(cgIdx[caller], callee)
+	}
+
+	// add the calls from planned functions
+	for _, f := range c.funcsCode {
+		fidx := c.funcs[f.name]
+		cgIdx[fidx] = findCallees(f.code.Func.Expr.Instrs)
+	}
+
+	keepFuncs := map[uint32]struct{}{}
+
+	// we'll keep
+	// - what's referenced in a table (these could be called indirectly)
+	// - what's exported or imported
+	// - what's been compiled by us
+	// - anything transitively called from those
+
+	for _, imp := range c.module.Import.Imports {
+		if _, ok := imp.Descriptor.(module.FunctionImport); ok {
+			reach(cgIdx, keepFuncs, c.funcs[imp.Name])
+		}
+	}
+
+	for _, exp := range c.module.Export.Exports {
+		if exp.Descriptor.Type == module.FunctionExportType {
+			reach(cgIdx, keepFuncs, c.funcs[exp.Name])
+		}
+	}
+
+	for _, f := range c.funcsCode {
+		reach(cgIdx, keepFuncs, c.funcs[f.name])
+	}
+
+	// anything referenced in a table
+	for _, seg := range c.module.Element.Segments {
+		for _, idx := range seg.Indices {
+			reach(cgIdx, keepFuncs, idx)
+		}
+	}
+
+	// remove all that's not needed, update index for remaining ones
+	funcNames := []module.NameMap{}
+	for _, nm := range c.module.Names.Functions {
+		if _, ok := keepFuncs[nm.Index]; ok {
+			funcNames = append(funcNames, nm)
+		}
+	}
+	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`,
+	// empty bodies would not.
+	nopEntry := module.Function{
+		Expr: module.Expr{
+			Instrs: []instruction.Instruction{instruction.Unreachable{}},
+		},
+	}
+	var buf bytes.Buffer
+	if err := encoding.WriteCodeEntry(&buf, &module.CodeEntry{Func: nopEntry}); err != nil {
+		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()
+		}
+	}
+	return nil
+}
+
+func findCallees(instrs []instruction.Instruction) []uint32 {
+	var ret []uint32
+	for _, expr := range instrs {
+		switch expr := expr.(type) {
+		case instruction.Call:
+			ret = append(ret, expr.Index)
+		case instruction.StructuredInstruction:
+			ret = append(ret, findCallees(expr.Instructions())...)
+		}
+	}
+	return ret
+}
+
+func reach(cg map[uint32][]uint32, keep map[uint32]struct{}, node uint32) {
+	if _, ok := keep[node]; !ok {
+		keep[node] = struct{}{}
+		for _, v := range cg[node] {
+			reach(cg, keep, v)
+		}
+	}
+}

internal/compiler/wasm/wasm.go

@@ -235,6 +235,9 @@ func New() *Compiler {
 		c.compileFuncs,
 		c.compilePlans,
 
+		// "local" optimizations
+		c.removeUnusedCode,
+
 		// final emissions
 		c.emitFuncs,
 
@@ -300,7 +303,16 @@ func (c *Compiler) initModule() error {
 
 	c.funcs = make(map[string]uint32)
 	for _, fn := range c.module.Names.Functions {
-		c.funcs[fn.Name] = fn.Index
+		name := fn.Name
+		// Account for recording duplicate functions -- this only happens
+		// with the RE2 C++ lib so far.
+		// NOTE: This isn't good enough for function names used more than
+		// two times. But let's deal with that when it happens.
+		if _, ok := c.funcs[name]; ok { // already seen
+			c.debug.Printf("function name duplicate: %s (%d)", name, fn.Index)
+			name = name + ".1"
+		}
+		c.funcs[name] = fn.Index
 	}
 
 	for _, fn := range c.policy.Funcs.Funcs {

main.go

@@ -22,4 +22,4 @@ func main() {
 //go:generate build/gen-run-go.sh internal/cmd/genopacapabilities/main.go capabilities.json
 
 // WASM base binary generation:
-//go:generate build/gen-run-go.sh internal/cmd/genopawasm/main.go -o internal/compiler/wasm/opa/opa.go internal/compiler/wasm/opa/opa.wasm
+//go:generate build/gen-run-go.sh internal/cmd/genopawasm/main.go -o internal/compiler/wasm/opa/opa.go internal/compiler/wasm/opa/opa.wasm  internal/compiler/wasm/opa/callgraph.csv

wasm/Makefile

@@ -67,7 +67,7 @@ push-builder:
 
 .PHONY: build
 build:
-	@$(DOCKER) run $(DOCKER_FLAGS) -v $(CURDIR):/src $(WASM_BUILDER_IMAGE) make $(WASM_OBJ_DIR)/opa.wasm
+	@$(DOCKER) run $(DOCKER_FLAGS) -v $(CURDIR):/src $(WASM_BUILDER_IMAGE) make $(WASM_OBJ_DIR)/opa.wasm $(WASM_OBJ_DIR)/callgraph.csv
 
 .PHONY: test
 test:
@@ -158,3 +158,9 @@ $(WASM_OBJ_DIR)/opa-test.wasm: $(OBJS) $(CPP_OBJS) $(LIB_OBJS) $(LIB_MPDEC_OBJS)
 			--import-memory \
 			--no-entry \
 			-o $@ $^
+
+$(WASM_OBJ_DIR)/callgraph.csv: $(WASM_OBJ_DIR)/opa.wasm
+	# NOTE: wasm-opt will output "warning: no output file specified",
+	# because we're not actually optimizing the wasm, but only extract
+	# information.
+	build/gen-wasm-callgraph.sh $< > $@

wasm/build/gen-wasm-callgraph.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -eo pipefail
+
+wasm-opt --print-call-graph $1 |
+  awk -F\" '/\/\/ call/{ print $2 "," $4 }'