perf: inline constant shifts and lower rotation idioms

pkg/test/zkc_unit_test.go

@@ -851,6 +851,10 @@ func Test_ZkcUnit_Shift_13(t *testing.T) {
 	checkZkcUnit(t, "zkc/unit/shift_13", DEFAULT_UNIT_CONFIG)
 }
 
+func Test_ZkcUnit_Rotate_01(t *testing.T) {
+	checkZkcUnit(t, "zkc/unit/rotate_01", DEFAULT_UNIT_CONFIG)
+}
+
 // ===================================================================
 // Read-Write Memory (RAM) Tests
 // ===================================================================

pkg/zkc/vm/internal/transform/constant_scan.go

@@ -0,0 +1,132 @@
+// Copyright Consensys Software Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+//
+// SPDX-License-Identifier: Apache-2.0
+package transform
+
+import (
+	"math/big"
+
+	"github.com/LFDT-Lineth/zkc/pkg/zkc/vm/internal/bytecode"
+	"github.com/LFDT-Lineth/zkc/pkg/zkc/vm/internal/descriptor"
+	"github.com/LFDT-Lineth/zkc/pkg/zkc/vm/internal/word"
+)
+
+// scanConstantRegisters identifies registers of a function which hold a
+// compile-time constant: registers with exactly one defining bytecode, where
+// that bytecode is a single-target ADD/SUB/MUL whose sources (if any) are all
+// themselves constant.  This covers plain constant loads (an ADD with no
+// sources) and simple constant arithmetic over them (e.g. "64 - n" where n was
+// loaded as a constant, as arises after inlining a call with literal
+// arguments).  A register whose folded value would not fit its declared
+// bitwidth is not reported, since executing the defining bytecode would fail
+// anyway.
+//
+// The result maps each such register to its value.  Registers written more
+// than once, written by any other bytecode form, or function inputs (written
+// by no bytecode) are never reported.
+func scanConstantRegisters[W word.Word[W]](fn *descriptor.Function[W]) map[bytecode.RegisterId]W {
+	var (
+		regs   = fn.Registers()
+		writes = make(map[bytecode.RegisterId]uint)
+		defOf  = make(map[bytecode.RegisterId]*bytecode.Arith[W])
+	)
+	// Record, for every register, its write count and (single-target Arith)
+	// defining bytecode.
+	for _, vec := range fn.Vectors() {
+		for _, insn := range vec.Bytecodes {
+			for _, target := range insn.Definitions() {
+				writes[target]++
+				//
+				if a, ok := insn.(*bytecode.Arith[W]); ok && len(a.Target) == 1 {
+					defOf[target] = a
+				} else {
+					defOf[target] = nil
+				}
+			}
+		}
+	}
+	// Iterate to a fixpoint, folding definitions whose sources are all known.
+	consts := make(map[bytecode.RegisterId]W)
+	//
+	for changed := true; changed; {
+		changed = false
+		//
+		for r, a := range defOf {
+			if _, done := consts[r]; done || a == nil || writes[r] != 1 {
+				continue
+			}
+			//
+			reg := regs[r]
+			if reg.IsNative() {
+				continue
+			}
+			//
+			value, ok := foldConstantArith(a, consts)
+			if !ok {
+				continue
+			}
+			// Reject values which do not fit the target register.
+			if value.Sign() < 0 || value.BitLen() > int(reg.Bitwidth().Unwrap()) {
+				continue
+			}
+			//
+			var w W
+
+			consts[r] = w.SetBigInt(value)
+			changed = true
+		}
+	}
+	//
+	return consts
+}
+
+// foldConstantArith evaluates a single-target ADD/SUB/MUL bytecode whose
+// sources are all in consts, following the Arith semantics
+// "target = sources[0] op sources[1] op ... op constant".
+func foldConstantArith[W word.Word[W]](a *bytecode.Arith[W],
+	consts map[bytecode.RegisterId]W,
+) (*big.Int, bool) {
+	switch a.Op {
+	case bytecode.OP_ADD, bytecode.OP_SUB, bytecode.OP_MUL:
+		// supported
+	default:
+		return nil, false
+	}
+	//
+	values := make([]*big.Int, 0, len(a.Source)+1)
+	//
+	for _, src := range a.Source {
+		w, ok := consts[src]
+		if !ok {
+			return nil, false
+		}
+		//
+		values = append(values, w.BigInt())
+	}
+	//
+	values = append(values, a.Constant.BigInt())
+	//
+	acc := new(big.Int).Set(values[0])
+	//
+	for _, v := range values[1:] {
+		switch a.Op {
+		case bytecode.OP_ADD:
+			acc.Add(acc, v)
+		case bytecode.OP_SUB:
+			acc.Sub(acc, v)
+		case bytecode.OP_MUL:
+			acc.Mul(acc, v)
+		}
+	}
+	//
+	return acc, true
+}

pkg/zkc/vm/internal/transform/lower_bitwise.go

@@ -38,8 +38,10 @@ func LowerBitwise[W word.Word[W]](program descriptor.Program[W]) descriptor.Prog
 
 	for i, mod := range out {
 		if fn, ok := mod.(*descriptor.Function[W]); ok {
+			consts := scanConstantRegisters(fn)
+			rots, deadShifts := scanRotations(fn, consts)
 			out[i] = lowerBitwiseFunction(fn, func(b Bytecode[W], alloc split.Allocator[W]) []Bytecode[W] {
-				return lowerBitwiseCode(b, alloc, helpers)
+				return lowerBitwiseCode(b, alloc, helpers, consts, rots, deadShifts)
 			})
 		}
 	}
@@ -77,29 +79,104 @@ func lowerBitwiseCode[W word.Word[W]](
 	b Bytecode[W],
 	registers split.Allocator[W],
 	helpers *shiftHelpers[W],
+	consts map[bytecode.RegisterId]W,
+	rots map[*bytecode.Bitwise[W]]rotation,
+	deadShifts map[*bytecode.Bitwise[W]]bool,
 ) []Bytecode[W] {
 	//
 	bw, ok := b.(*bytecode.Bitwise[W])
 	if !ok {
 		return []Bytecode[W]{b}
 	}
+	// Shifts absorbed into a rotation (see scanRotations) are dropped: the OR
+	// which consumed them is rewritten below.
+	if deadShifts[bw] {
+		return nil
+	}
 	//
 	switch bw.Op {
 	case bytecode.OP_NOT:
 		return inlineBitwiseNot(bw, registers)
 	case bytecode.OP_SHL, bytecode.OP_SHR:
-		return lowerBitwiseShlShr(bw, registers, helpers)
+		return lowerBitwiseShlShr(bw, registers, helpers, consts)
+	case bytecode.OP_OR:
+		if rot, ok := rots[bw]; ok {
+			return lowerRotation(bw, registers, helpers, rot)
+		}
+		// Otherwise lowered after register splitting; see LowerOrXorAnd.
+		return []Bytecode[W]{b}
 	default:
-		// AND/OR/XOR are lowered after register splitting; see LowerOrXorAnd.
+		// AND/XOR are lowered after register splitting; see LowerOrXorAnd.
 		return []Bytecode[W]{b}
 	}
 }
 
+// lowerRotation rewrites an OR recognised as a rotation idiom: a constant
+// rotation is realised inline (Destruct + swapped Concat, no helper), while a
+// dynamic one becomes a CALL into the rotation chain (see ensureRot), whose
+// entry level is the amount register's width.
+func lowerRotation[W word.Word[W]](
+	b *bytecode.Bitwise[W],
+	registers split.Allocator[W],
+	helpers *shiftHelpers[W],
+	rot rotation,
+) []Bytecode[W] {
+	width := uint(b.Bitwidth)
+	//
+	if rot.isConst {
+		return inlineRotlByConst(b.Target, rot.source, width, rot.constAmount, registers)
+	}
+	//
+	var (
+		amtWidth = registers.Registers()[rot.amount].Bitwidth().Unwrap()
+		id       = helpers.ensureRot(rot.op, width, amtWidth)
+	)
+	//
+	return []Bytecode[W]{
+		bytecode.CallFun[W](uint16(id), []bytecode.RegisterId{rot.source, rot.amount}, []bytecode.RegisterId{b.Target}),
+	}
+}
+
+// inlineRotlByConst emits "target = source rotl s" for a compile-time
+// constant amount s in [0, width] directly into the caller's bytecode
+// stream: both bounds are a move, anything else is a Destruct into
+// [lo:u(width-s), hi:u_s] followed by the swapped Concat target = lo : hi
+// (hi in the low bits) — no field arithmetic, no helper modules, no lookups.
+func inlineRotlByConst[W word.Word[W]](target, source bytecode.RegisterId,
+	width, s uint, registers split.Allocator[W],
+) []Bytecode[W] {
+	var zero = word.Const64[W](0)
+	//
+	s %= width
+	//
+	if s == 0 {
+		return []Bytecode[W]{bytecode.AddConst(target, []bytecode.RegisterId{source}, zero)}
+	}
+	//
+	var (
+		lo = registers.Allocate("", util.Some(width-s))
+		hi = registers.Allocate("", util.Some(s))
+	)
+	//
+	return []Bytecode[W]{
+		bytecode.AddVec[W]([]bytecode.RegisterId{lo, hi}, []bytecode.RegisterId{source}),
+		bytecode.AssignV[W]([]bytecode.RegisterId{target}, hi, lo),
+	}
+}
+
 func lowerBitwiseShlShr[W word.Word[W]](
 	b *bytecode.Bitwise[W],
 	registers split.Allocator[W],
 	helpers *shiftHelpers[W],
+	consts map[bytecode.RegisterId]W,
 ) []Bytecode[W] {
+	// Fast path: a shift by a compile-time constant needs no barrel chain at
+	// all — it is realised inline by the same Destruct / Concat scheme the
+	// chain levels use internally (see shiftByConst).
+	if amount, ok := consts[b.Right.AsRegister()]; ok {
+		return inlineShiftByConst(b, registers, amount)
+	}
+	//
 	var (
 		// NOTE: the shift amount is always a register (constant operands are
 		// only supported for AND/OR/XOR).
@@ -115,6 +192,55 @@ func lowerBitwiseShlShr[W word.Word[W]](
 	}
 }
 
+// inlineShiftByConst emits "target = left op amount" for a compile-time
+// constant amount directly into the caller's bytecode stream, mirroring
+// shiftByConst: a shift of zero is a move, an amount >= width yields zero, and
+// anything else is a Destruct (SHR) or Destruct + Concat (SHL) — no field
+// arithmetic, no helper modules, no lookups.
+func inlineShiftByConst[W word.Word[W]](b *bytecode.Bitwise[W],
+	registers split.Allocator[W], amount W,
+) []Bytecode[W] {
+	var (
+		width = uint(b.Bitwidth)
+		zero  = word.Const64[W](0)
+	)
+	// Amounts >= width shift everything out.
+	if amount.Cmp64(uint64(width)) >= 0 {
+		return []Bytecode[W]{bytecode.LoadConst(b.Target, zero)}
+	}
+	//
+	shift := uint(amount.Uint64())
+	//
+	if shift == 0 {
+		return []Bytecode[W]{bytecode.AddConst(b.Target, []bytecode.RegisterId{b.Left}, zero)}
+	}
+	//
+	var (
+		drop = registers.Allocate("", util.Some(shift))
+		keep = registers.Allocate("", util.Some(width-shift))
+	)
+	//
+	switch b.Op {
+	case bytecode.OP_SHR:
+		// Destruct left into [drop, keep] (little-endian): keep = left >> shift.
+		return []Bytecode[W]{
+			bytecode.AddVec[W]([]bytecode.RegisterId{drop, keep}, []bytecode.RegisterId{b.Left}),
+			bytecode.AddConst(b.Target, []bytecode.RegisterId{keep}, zero),
+		}
+	case bytecode.OP_SHL:
+		zeros := registers.Allocate("", util.Some(shift))
+		// Destruct left into [keep, drop] (little-endian): keep = left mod
+		// 2^(width-shift), then target = keep : zeros.
+		return []Bytecode[W]{
+			bytecode.AddVec[W]([]bytecode.RegisterId{keep, drop}, []bytecode.RegisterId{b.Left}),
+			bytecode.LoadConst(zeros, zero),
+			bytecode.AssignV[W]([]bytecode.RegisterId{b.Target}, zeros, keep),
+		}
+	default:
+		panic("expected shift operation")
+	}
+}
+
 // inlineBitwiseNot emits ~x as (MASK - x) directly into the caller's bytecode
 // stream, where MASK = 2^width - 1.  No helper module is created.
 func inlineBitwiseNot[W word.Word[W]](b *bytecode.Bitwise[W], registers split.Allocator[W]) []Bytecode[W] {

pkg/zkc/vm/internal/transform/lower_shift.go

@@ -84,6 +84,9 @@ type shiftHelperKey struct {
 	// width seen across its call sites (always > the chain depth, so guard and
 	// level keys never collide).
 	amtWidth uint
+	// rot separates rotation helpers (see ensureRot) from shift helpers: a
+	// rotation chain reuses op to select its direction (OP_SHL: left).
+	rot bool
 }
 
 // shiftHelpers is the registry of SHL/SHR helper modules built by