perf: ir peephole pass

compiler-architecture.md

@@ -456,9 +456,6 @@ check_enough_funds($pulled_sn, $share_n)
 
 ### Optimisations
 
-> [!NOTE]
-> Peephole optimisations aren't implemented yet. There is a POC in another branch, to measure how much perf could be impacted, but it's too soon to consider
-
 You may have noticed that the previous compilation examples emit _a lot_ of garbage.
 That's done by design: the compiler must be simple and declarative. We don't want dozens of special cases in the compilation logic, which must express a general, albeit redundant, template which focuses on correctness.
 
@@ -468,29 +465,57 @@ That change also *enables* a peephole that was previously out of reach. `assert_
 
 However, the `irInstr` layer allows us to easily rewrite the instructions so that we remove garbage instructions, precomputing more aggressively, rewrite them into more efficient code (this is called [peephole optimisation](https://en.wikipedia.org/wiki/Peephole_optimization))
 
-Each peephole is independent and is expressed as a `func(instr []irInstr) []irInstr`, which returns the new instructions set, or nil if it didn't change.
-Each peephole is independently testable and reviewable.
+Each peephole is independent and is expressed as an `ir.Peephole`, a `func([]ir.Instr) []ir.Instr` which returns the new instruction stream, or nil if it didn't change anything. A pass never mutates its input.
+Each peephole is independently testable and reviewable: one file per pass under `internal/ir/peephole_*.go`, with its tests written against the textual IR.
+
+We apply each peephole optimisation sequentially, and repeat until a full sweep changes nothing (the program `p` such that `f(p) == p` for every peephole `f`).
+
+Note that proving that a peephole function _does_ have a fixed point is usually simple, whereas proving that the function composition of all the peepholes isn't. Pratically speaking, we avoid a non-terminating optimisation by imposing a max amount of sweeps (`maxOptimizeSweeps`). A clever order of peepholes makes convergence fast anyway: `ir.DefaultPeepholes()` is ordered so that one sweep normally does it all and the second one confirms the fixed point.
+
+Optimisation is **opt-in**. `compiler.Compile` assembles the plain stream; `compiler.CompileWithOptions(program, flags, compiler.Options{Optimize: true})` — `numscript.CompileOptions{Optimize: true}` on the public API — runs `ir.Optimize` first, and typechecks the result again before assembling it. The host decides, per compilation.
+
+The contract of every pass is that the optimized program is **behaviorally identical** to the plain one: same postings in the same order, same metadata, same errors — including which error surfaces when more than one could. Three things enforce it:
+
+- the spec corpus (`TestCompilerScripts`) runs the plain and the optimized program on every case and requires identical results and identical error structs, before checking the plain one against the spec;
+- `TestCompiledCorpusPassesVerify` runs both through `vm.Verify` and the `Encode`/`Decode` round trip;
+- the differential sweep (`internal/difftest`) has a `vm vs optimized vm` leg compared with `CompareIdentical`: no aggregation, and one tolerance, a capacity bound of the plain program the optimized one does not hit.
 
-We apply each peephole optimisation sequentially, and repeat until we reach a fixed point for each peephole (the program `p` such that `f(p) == p`, where `f` is the peephole function).
+#### Implemented peepholes
 
-Note that proving that a peephole function _does_ have a fixed point is usually simple, whereas proving that the function composition of all the peepholes isn't. Pratically speaking, we can avoid non-terminating optimisation passes by imposing a max amount of optimisation passes. A clever order of peepholes should make convergence quite fast anyway.
+Three of them share one analysis, `walkConsts`: a forward walk that knows, before each instruction, which registers hold the same constant on every path reaching it. Jumps only go forward, so every edge into a label is known when the walk reaches it and one pass is exact. Unlike a "written once" rule, it does not depend on the typechecker, which checks reads against writes in program order, not per path.
+
+1. **`constProp`** — a pure op whose result is the same int, str or bool on every path reaching it becomes a load of that value. Portions are tracked but not materialized, as there is no portion load. Each op is folded exactly as the VM computes it; a `mk_portion` with a zero denominator is left to fail at run time. `check_enough_funds` fails unless `got == needed`, so after it `got` takes `needed`'s value: with a literal amount, the whole allotment split is computed at compile time.
+2. **`dropProvenAsserts`** — a check that passes on every path reaching it goes: the `assert_*` family and `check_enough_funds`. A check that may fail stays where it is, so the first error a run hits is unchanged. Dropping `assert_leftover` is what frees the portion arithmetic for `deadCode`.
+3. **`pruneDeadBranches`** — a conditional jump on a constant bool becomes an unconditional `jmp` or nothing; the instructions between a `jmp` and the next label a reachable jump targets go; a jump to the next instruction goes; a label no jump targets goes.
+4. **`deadCode`** — a pure instruction whose dests no path reads goes. Store reads, pulls and every check stay; a `mk_portion` goes only where its denominator is known not to be zero.
+
+The `@world` diamond (see [Plain account source/dest](#plain-account-sourcedest-bounded)) is what the folding and the pruning are for: `str_eq` over two `load_str` folds, the branch on it goes, and a literal `@world` source is back to a single unbounded `pull_account` — any other literal account to a single bounded one. The amount clamp (`max($amount, 0)`) collapses the same way when the amount is a literal.
+
+`send [USD/2 42] (source = @world destination = @dest)` optimizes to:
+
+```
+  $r0 = "USD/2"
+  $r1 = 42
+  set_current_asset($r0)
+  $r2 = 42
+  $r3 = "world"
+  $r4 = pull_account(account: $r3, cap: $r2)
+  check_enough_funds($r4, $r1)
+  $r5 = "dest"
+  send_to_account(account: $r5)
+```
 
-> TODO list some peepholes
+and `send [USD/2 100] (source = @src destination = {1/3 to @a 1/3 to @b remaining to @c})` to three sends of the constants 34, 33 and 33: no portion op is left.
 
-The `@world` diamond (see [Plain account source/dest](#plain-account-sourcedest-bounded)) is
-the clearest case to date, and it needs two passes that compose:
+The passes can trade one capacity bound of the assembler for another (a shorter stream with longer-lived registers), so when the optimized stream hits `ErrRegisterBankOverflow` or `ErrProgramTooLarge` the compiler assembles the plain one instead. Optimizing can lift a bound (`deadCode` drops unused `load_var`s, so a script declaring more vars than a bank holds may compile), never impose one.
 
-1. **Const-fold `str_eq`** — when both operands trace back to `load_str` constants, the
-   comparison is statically decidable: replace it with `true` or `false`.
-2. **Dead-branch elimination** — a conditional jump on a register holding a known bool
-   is either a no-op or an unconditional `jmp`; then everything between a `jmp` and the
-   next reachable label is unreachable. `is_zero` over a `load_int` folds the same way,
-   which is what makes the quantity branches reachable for this pass too.
+#### Planned peepholes
 
-Together they collapse a literal `@world` source back to a single unbounded
-`pull_account`, and any other literal account back to a single bounded one — i.e. to
-exactly the code the compiler emitted when the VM still special-cased the name. The
-prologue's `load_str "world"` also becomes dead once no branch refers to it.
+- **Post from unbounded** — a send statement with a single unbounded source (`@world`, or `allowing unbounded overdraft`) and a single destination becomes a direct posting, skipping the funds queue.
+- **Dead credit** — the destination credit of a direct posting whose destination's balance is never read.
+- **Funds bypass** — the Pull-then-Send → take + post rewrite for a bounded single source, a fan-out (one source, allotment destinations) and a fan-in (statically distinct sources, one destination). Never N×M, `oneof` or scaling.
+- **Integer allotment split** — `portion_to_int(mul_portion(p, int_to_portion(x)))` as one integer op, so an allotment over a var amount allocates no `big.Rat`. Needs a new opcode.
+- **Register allocation** — only if the benchmarks justify it; see below.
 
 ### Registers allocator
 

internal/compiler/bench_test.go

@@ -260,12 +260,23 @@ func BenchmarkCompiledVMCapped(b *testing.B) {
 // reuse one Vm, measure only the run.
 func benchCompiledVM(b *testing.B, src string, store benchStore) {
 	b.Helper()
+	benchCompiledVMWith(b, src, store, compiler.Options{})
+}
+
+// benchCompiledVMOpt is benchCompiledVM on the optimized program.
+func benchCompiledVMOpt(b *testing.B, src string, store benchStore) {
+	b.Helper()
+	benchCompiledVMWith(b, src, store, compiler.Options{Optimize: true})
+}
+
+func benchCompiledVMWith(b *testing.B, src string, store benchStore, opts compiler.Options) {
+	b.Helper()
 
 	parsed := parser.Parse(src)
 	if len(parsed.Errors) != 0 {
 		b.Fatalf("parse errors: %v", parsed.Errors)
 	}
-	_, program, err := compiler.Compile(parsed.Value, nil)
+	_, program, err := compiler.CompileWithOptions(parsed.Value, nil, opts)
 	if err != nil {
 		b.Fatalf("compile: %v", err)
 	}
@@ -315,6 +326,38 @@ func BenchmarkCompiledVMAllotmentThirds(b *testing.B) {
 	}})
 }
 
+// --- Optimized programs ----------------------------------------------------
+// The same scripts through compiler.Options{Optimize: true}, so each peephole's
+// effect is a before/after on the same benchmark.
+
+func simpleStore() benchStore {
+	return benchStore{balances: map[funds.PairKey]*big.Int{
+		{Account: "src", Asset: "USD/2", Color: ""}: big.NewInt(100),
+	}}
+}
+
+func BenchmarkCompiledVMOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrc, simpleStore())
+}
+
+func BenchmarkCompiledVMCappedOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcCapped, cappedStore())
+}
+
+// The most common numscript shape: funding an account from @world.
+const benchSrcWorld = `send [USD/2 42] (
+	source = @world
+	destination = @dest
+)`
+
+func BenchmarkCompiledVMWorld(b *testing.B) {
+	benchCompiledVM(b, benchSrcWorld, benchStore{})
+}
+
+func BenchmarkCompiledVMWorldOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcWorld, benchStore{})
+}
+
 // Fan-in: {1/3 from @a; 1/3 from @b; 1/3 from @c} -> @dest. Allotment on the
 // source side, no early-exit jump.
 const benchSrcFanIn = `send [USD/2 30] (
@@ -333,3 +376,72 @@ func BenchmarkCompiledVMFanIn(b *testing.B) {
 		{Account: "c", Asset: "USD/2", Color: ""}: big.NewInt(100),
 	}})
 }
+
+func BenchmarkCompiledVMAllotmentOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcAllotment, benchStore{balances: map[funds.PairKey]*big.Int{
+		{Account: "src", Asset: "USD/2", Color: ""}: big.NewInt(1000),
+	}})
+}
+
+func BenchmarkCompiledVMAllotmentThirdsOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcAllotmentThirds, benchStore{balances: map[funds.PairKey]*big.Int{
+		{Account: "src", Asset: "USD/2", Color: ""}: big.NewInt(1000),
+	}})
+}
+
+func BenchmarkCompiledVMFanInOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcFanIn, benchStore{balances: map[funds.PairKey]*big.Int{
+		{Account: "a", Asset: "USD/2", Color: ""}: big.NewInt(100),
+		{Account: "b", Asset: "USD/2", Color: ""}: big.NewInt(100),
+		{Account: "c", Asset: "USD/2", Color: ""}: big.NewInt(100),
+	}})
+}
+
+// The same split with the amount from a var: no constant reaches the portion
+// arithmetic, so it runs on every execution.
+const benchSrcAllotmentVar = `vars { monetary $m }
+send $m (
+	source = @src
+	destination = {
+		10% to @fees
+		remaining to @b
+	}
+)`
+
+func benchVarsWith(b *testing.B, opts compiler.Options) {
+	parsed := parser.Parse(benchSrcAllotmentVar)
+	enc, program, err := compiler.CompileWithOptions(parsed.Value, nil, opts)
+	if err != nil {
+		b.Fatal(err)
+	}
+	vars, err := enc.Encode(map[string]string{"m": "USD/2 100"})
+	if err != nil {
+		b.Fatal(err)
+	}
+	store := benchStore{balances: map[funds.PairKey]*big.Int{{Account: "src", Asset: "USD/2"}: big.NewInt(1000)}}
+	machine := vm.NewVm(program)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := vm.Exec(context.Background(), machine, &vars, store); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCompiledVMAllotmentVar(b *testing.B) { benchVarsWith(b, compiler.Options{}) }
+func BenchmarkCompiledVMAllotmentVarOpt(b *testing.B) {
+	benchVarsWith(b, compiler.Options{Optimize: true})
+}
+
+const benchSrcAllotmentPct = `send [USD/2 100] (
+	source = @src
+	destination = {
+		10% to @fees
+		remaining to @b
+	}
+)`
+
+func BenchmarkCompiledVMAllotmentPctOpt(b *testing.B) {
+	benchCompiledVMOpt(b, benchSrcAllotmentPct, benchStore{balances: map[funds.PairKey]*big.Int{{Account: "src", Asset: "USD/2"}: big.NewInt(1000)}})
+}

internal/compiler/compiler.go

@@ -1,6 +1,8 @@
 package compiler
 
 import (
+	"errors"
+	"fmt"
 	"maps"
 	"math/big"
 	"slices"
@@ -21,6 +23,18 @@ import (
 // interpreter.RunProgram, the script's own #![feature(..)] declarations are
 // unioned in.
 func Compile(program parser.Program, featureFlags map[string]struct{}) (VarsEncoder, vm.Program, error) {
+	return CompileWithOptions(program, featureFlags, Options{})
+}
+
+// Options configures CompileWithOptions. The zero value is what Compile uses.
+type Options struct {
+	// Optimize runs ir.Optimize over the instruction stream before assembling.
+	// The optimized program produces the same postings, metadata and errors as
+	// the plain one.
+	Optimize bool
+}
+
+func CompileWithOptions(program parser.Program, featureFlags map[string]struct{}, opts Options) (VarsEncoder, vm.Program, error) {
 	compiled, cErr := compileProgramToIR(program, featureFlags)
 	if cErr != nil {
 		return VarsEncoder{}, vm.Program{}, cErr
@@ -30,7 +44,21 @@ func Compile(program parser.Program, featureFlags map[string]struct{}) (VarsEnco
 		return VarsEncoder{}, vm.Program{}, err
 	}
 
-	prog, err := ir.Assemble(compiled.instructions)
+	instrs := compiled.instructions
+	if opts.Optimize {
+		instrs = ir.Optimize(instrs)
+		// a pass that breaks register typing is a bug in the pass
+		if err := ir.Typecheck(instrs); err != nil {
+			return VarsEncoder{}, vm.Program{}, fmt.Errorf("optimized program: %w", err)
+		}
+	}
+
+	prog, err := ir.Assemble(instrs)
+	// the passes can trade one capacity bound for another (a shorter stream
+	// with longer-lived registers), so the plain stream decides those errors
+	if opts.Optimize && (errors.Is(err, ir.ErrRegisterBankOverflow) || errors.Is(err, ir.ErrProgramTooLarge)) {
+		prog, err = ir.Assemble(compiled.instructions)
+	}
 	if err != nil {
 		return VarsEncoder{}, vm.Program{}, err
 	}

internal/compiler/optimize_test.go

@@ -0,0 +1,125 @@
+package compiler
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/formancehq/numscript/internal/ir"
+	"github.com/formancehq/numscript/internal/parser"
+	"github.com/gkampitakis/go-snaps/snaps"
+	"github.com/stretchr/testify/require"
+)
+
+// getOptimizedOutput is getCompiledOutput after ir.Optimize. The dump is
+// canonicalized through a parse: a pass that deletes an instruction leaves a
+// gap in the register numbering, and the parser renumbers by first appearance.
+func getOptimizedOutput(t *testing.T, source string) string {
+	t.Helper()
+	program := parser.Parse(source)
+	require.Empty(t, program.Errors)
+	compiled, err := compileProgramToIR(program.Value, nil)
+	require.Nil(t, err)
+	require.NoError(t, ir.Typecheck(compiled.instructions))
+
+	optimized := ir.Optimize(compiled.instructions)
+	require.NoError(t, ir.Typecheck(optimized))
+
+	instrs, errs := ir.Parse(ir.Dump(optimized))
+	require.Empty(t, errs, "the optimized dump does not parse back")
+	return "\n" + ir.Dump(instrs)
+}
+
+// A literal @world keeps only the unbounded arm of its branch.
+func TestOptimizedWorldSend(t *testing.T) {
+	out := getOptimizedOutput(t, `
+		send [USD/2 42] (
+			source = @world
+			destination = @dest
+		)
+	`)
+
+	snaps.MatchInlineSnapshot(t, out, snaps.Inline(`
+  $r0 = "USD/2"
+  $r1 = 42
+  set_current_asset($r0)
+  $r2 = 42
+  $r3 = "world"
+  $r4 = pull_account(account: $r3, cap: $r2)
+  check_enough_funds($r4, $r1)
+  $r5 = "dest"
+  send_to_account(account: $r5)
+`))
+}
+
+// Any other literal account keeps only the bounded one.
+func TestOptimizedBoundedSend(t *testing.T) {
+	out := getOptimizedOutput(t, `
+		send [USD/2 10] (
+			source = @src
+			destination = @dest
+		)
+	`)
+
+	snaps.MatchInlineSnapshot(t, out, snaps.Inline(`
+  $r0 = "USD/2"
+  $r1 = 10
+  set_current_asset($r0)
+  $r2 = 10
+  $r3 = "src"
+  $r4 = 0
+  $r5 = pull_account(account: $r3, cap: $r2, overdraft: $r4)
+  check_enough_funds($r5, $r1)
+  $r6 = "dest"
+  send_to_account(account: $r6)
+`))
+}
+
+// An account from a var is decided at run time: both arms stay.
+func TestOptimizedSendFromVarKeepsBothArms(t *testing.T) {
+	out := getOptimizedOutput(t, `
+		vars { account $src }
+		send [USD/2 10] (
+			source = $src
+			destination = @dest
+		)
+	`)
+
+	snaps.MatchInlineSnapshot(t, out, snaps.Inline(`
+  $r0 = "world"
+  $r1 = load_var<str>(0)
+  $r2 = "USD/2"
+  $r3 = 10
+  set_current_asset($r2)
+  $r4 = 10
+  $r5 = 0
+  $r6 = str_eq($r1, $r0)
+  jmp_if_false($r6, #not_world_1)
+  $r7 = pull_account(account: $r1, cap: $r4)
+  jmp(#pull_end_2)
+#not_world_1
+  $r7 = pull_account(account: $r1, cap: $r4, overdraft: $r5)
+#pull_end_2
+  check_enough_funds($r7, $r3)
+  $r8 = "dest"
+  send_to_account(account: $r8)
+`))
+}
+
+// The unbounded arm the compiler emits for a literal @world is reached by a
+// branch on str_eq; folding and pruning it must leave a single pull. Checked on
+// the plain dump too, so a change in how the compiler spells the branch shows
+// up here rather than as a silently unoptimized program.
+func TestOptimizeCollapsesLiteralWorldBranch(t *testing.T) {
+	src := `send [USD/2 42] (source = @world destination = @dest)`
+
+	plain := getCompiledOutput(t, src)
+	require.Equal(t, 2, strings.Count(plain, "pull_account("), "the plain program has both arms")
+	require.Contains(t, plain, "str_eq(")
+
+	out := getOptimizedOutput(t, src)
+	require.Equal(t, 1, strings.Count(out, "pull_account("))
+	require.NotContains(t, out, "overdraft:", "the bounded arm is the dead one")
+	require.NotContains(t, out, "str_eq(")
+	require.NotContains(t, out, "jmp")
+	require.NotContains(t, out, "#")
+}