diff --git a/pkg/test/zkc_unit_test.go b/pkg/test/zkc_unit_test.go index 391e63552..c6580de99 100644 --- a/pkg/test/zkc_unit_test.go +++ b/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 // =================================================================== diff --git a/pkg/zkc/vm/internal/transform/constant_scan.go b/pkg/zkc/vm/internal/transform/constant_scan.go new file mode 100644 index 000000000..3b851af86 --- /dev/null +++ b/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 +} diff --git a/pkg/zkc/vm/internal/transform/lower_bitwise.go b/pkg/zkc/vm/internal/transform/lower_bitwise.go index 0dc74d30d..a24bcdc98 100644 --- a/pkg/zkc/vm/internal/transform/lower_bitwise.go +++ b/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] { diff --git a/pkg/zkc/vm/internal/transform/lower_shift.go b/pkg/zkc/vm/internal/transform/lower_shift.go index 2ac88028e..d6540c527 100644 --- a/pkg/zkc/vm/internal/transform/lower_shift.go +++ b/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 diff --git a/pkg/zkc/vm/internal/transform/rot_helpers.go b/pkg/zkc/vm/internal/transform/rot_helpers.go new file mode 100644 index 000000000..9f1c5ecd4 --- /dev/null +++ b/pkg/zkc/vm/internal/transform/rot_helpers.go @@ -0,0 +1,145 @@ +// 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 ( + "fmt" + + "github.com/LFDT-Lineth/zkc/pkg/schema/register" + "github.com/LFDT-Lineth/zkc/pkg/util" + "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" +) + +// ensureRot returns the module id of the rotation helper for values of the +// given width, entered at level amtWidth of the chain, creating any missing +// levels. op selects the direction: OP_SHL rotates left, OP_SHR right. +// Unlike shifts, rotations need no guard: callers guarantee amtWidth <= +// shiftChainDepth(width), and for the power-of-two widths recognised by +// scanRotations every amount value is then in range. +func (p *shiftHelpers[W]) ensureRot(op bytecode.Operation, width uint, amtWidth uint) uint { + key := shiftHelperKey{op: op, width: width, amtWidth: amtWidth, rot: true} + + if id, ok := p.ids[key]; ok { + return id + } + + var subID uint + if amtWidth > 1 { + subID = p.ensureRot(op, width, amtWidth-1) + } + + id := p.baseID + uint(len(p.items)) + p.ids[key] = id + p.items = append(p.items, newRotLevelHelper[W](key, subID)) + + return id +} + +// rotHelperName is the module name of a rotation helper: the direction, the +// value width and the amount (arg2) width. +func rotHelperName(key shiftHelperKey) string { + dir := "rotl" + if key.op == bytecode.OP_SHR { + dir = "rotr" + } + + return fmt.Sprintf("$bit_%s_u%d_u%d", dir, key.width, key.amtWidth) +} + +// newRotLevelHelper builds level j (= key.amtWidth) of the rotation chain for +// values of width w (= key.width), mirroring newShiftLevelHelper: +// +// let bit:u1, low:u(j-1) = n +// +// level_j(a, n:u_j) = level_{j-1}(bit == 0 ? a : a rot 2^(j-1), nlow) +// level_1(a, n:u1) = n == 0 ? a : a rot 1 +// +// where "a rot k" is a rotation by the constant k, realised purely by a +// Destruct followed by a swapped Concat (see rotByConst) — cheaper than a +// constant shift, which additionally zero-fills. subID is the module id of +// level j-1; it is ignored when j == 1. +func newRotLevelHelper[W word.Word[W]](key shiftHelperKey, subID uint) descriptor.Module[W] { + var padding W + + b := newHelperBuilder[W](key.width, 2) + b.base[1] = descriptor.NewRegister(register.INPUT_REGISTER, "arg2", util.Some(key.amtWidth), padding) + + a, n, out := b.inputs[0], b.inputs[1], b.output + width := key.width + level := key.amtWidth + zero := word.Const64[W](0) + + if level == 1 { + // if n == 0: return a + b.emit(bytecode.NewSkipIf(bytecode.CONDITION_NEQ, 2, + bytecode.NewRegisterVector(n), + bytecode.NewConstantOperand(zero))) + b.emit(bytecode.AddConst(out, []bytecode.RegisterId{a}, zero)) + b.emit(bytecode.NewRet[W]()) + // out = a rotated by 1 + b.emitAll(rotByConst(b, key.op, out, a, 1)) + b.emit(bytecode.NewRet[W]()) + } else { + shift := uint(1) << (level - 1) + // Destruct n into [nlow:u(level-1), bit:u1] (little-endian). + nlow := b.newComputedWidth("$nlow", level-1) + bit := b.newComputedWidth("$bit", 1) + b.emit(bytecode.AddVec[W]([]bytecode.RegisterId{nlow, bit}, []bytecode.RegisterId{n})) + // next = bit == 0 ? a : a rotated by 2^(level-1) + next := b.newComputedWidth("$next", width) + rotated := rotByConst(b, key.op, next, a, shift) + b.emit(bytecode.NewSkipIf(bytecode.CONDITION_NEQ, 2, + bytecode.NewRegisterVector(bit), + bytecode.NewConstantOperand(zero))) + b.emit(bytecode.AddConst(next, []bytecode.RegisterId{a}, zero)) + b.emit(bytecode.NewSkip[W](uint16(len(rotated)))) + b.emitAll(rotated) + // return level_{j-1}(next, nlow) + b.emit(bytecode.CallFun[W](uint16(subID), []bytecode.RegisterId{next, nlow}, []bytecode.RegisterId{out})) + b.emit(bytecode.NewRet[W]()) + } + + return descriptor.NewFunction(rotHelperName(key), b.regs(), descriptor.BYTECODE_FUNCTION, nil, + []BytecodeVector[W]{bytecode.NewVector(b.code...)}) +} + +// rotByConst returns the codes computing "target = a rot shift" for a +// constant amount in (0, width), allocating temporaries on the builder but +// NOT emitting (so the caller can size a Skip over the sequence). A +// rotate-left by s splits a into [lo:u(w-s), hi:u_s] (little-endian) and +// reassembles target = lo : hi with hi in the low bits; a rotate-right by +// shift is a rotate-left by width - shift. +func rotByConst[W word.Word[W]](b *helperBuilder[W], op bytecode.Operation, + target, a bytecode.RegisterId, shift uint, +) []Bytecode[W] { + var ( + width = b.width + s = shift + ) + + if op == bytecode.OP_SHR { + s = width - shift + } + + var ( + lo = b.newComputedWidth("$lo", width-s) + hi = b.newComputedWidth("$hi", s) + ) + + return []Bytecode[W]{ + bytecode.AddVec[W]([]bytecode.RegisterId{lo, hi}, []bytecode.RegisterId{a}), + bytecode.AssignV[W]([]bytecode.RegisterId{target}, hi, lo), + } +} diff --git a/pkg/zkc/vm/internal/transform/rotation_scan.go b/pkg/zkc/vm/internal/transform/rotation_scan.go new file mode 100644 index 000000000..8f746fd03 --- /dev/null +++ b/pkg/zkc/vm/internal/transform/rotation_scan.go @@ -0,0 +1,217 @@ +// 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/bits" + + "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" +) + +// rotation describes a recognised rotation idiom "(x << s) | (x >> t)" where +// s + t == width: an OR of two opposite shifts of the same source register. +// When both amounts are compile-time constants the rotation is realised +// inline as a single Destruct + Concat (see inlineRotlByConst); otherwise +// `amount` names the dynamic amount register and `op` selects the direction +// (OP_SHL: rotate left by amount; OP_SHR: rotate right by amount) — the +// other amount register is structurally "width - amount" and is discarded. +type rotation struct { + // source is the register being rotated. + source bytecode.RegisterId + // constAmount is the rotate-left amount, valid iff isConst. It lies in + // [0, width]; both bounds denote the identity rotation. + constAmount uint + isConst bool + // amount is the dynamic amount register (unused when isConst). + amount bytecode.RegisterId + // op is OP_SHL for rotate-left, OP_SHR for rotate-right (unused when + // isConst). + op bytecode.Operation +} + +// scanRotations identifies rotation idioms in a function: OR bytecodes whose +// two operands are single-use, single-writer results of a SHL and a SHR of +// the same (single-writer) source, with amounts summing to the operand +// width. Amounts are constant when both appear in consts; a dynamic pair +// qualifies when one amount register is structurally defined as "width - +// other" (a single-writer two-source SUB whose first source is the constant +// width) — the shape produced by the usual "(x << n) | (x >> (w - n))" +// source-level rotation. Dynamic recognition is limited to power-of-two +// widths, where the amount register (of width <= log2(width)) always holds +// an in-range rotation. +// +// The result maps each such OR to its rotation, alongside the set of shift +// bytecodes it obsoletes (which the caller must drop). +func scanRotations[W word.Word[W]](fn *descriptor.Function[W], + consts map[bytecode.RegisterId]W, +) (map[*bytecode.Bitwise[W]]rotation, map[*bytecode.Bitwise[W]]bool) { + var ( + regs = fn.Registers() + writes = make(map[bytecode.RegisterId]uint) + uses = make(map[bytecode.RegisterId]uint) + defOf = make(map[bytecode.RegisterId]Bytecode[W]) + rots = make(map[*bytecode.Bitwise[W]]rotation) + dead = make(map[*bytecode.Bitwise[W]]bool) + ) + // + for _, vec := range fn.Vectors() { + for _, insn := range vec.Bytecodes { + for _, target := range insn.Definitions() { + writes[target]++ + defOf[target] = insn + } + // + for _, src := range insn.Uses() { + uses[src]++ + } + } + } + // + for _, vec := range fn.Vectors() { + for _, insn := range vec.Bytecodes { + or, ok := insn.(*bytecode.Bitwise[W]) + if !ok || or.Op != bytecode.OP_OR || or.Right.IsConstant() { + continue + } + // + right := or.Right.AsRegister() + if or.Left == right { + continue + } + // Both operands must be single-writer shifts consumed only here. + if writes[or.Left] != 1 || writes[right] != 1 || uses[or.Left] != 1 || uses[right] != 1 { + continue + } + // + shl, okL := defOf[or.Left].(*bytecode.Bitwise[W]) + shr, okR := defOf[right].(*bytecode.Bitwise[W]) + // + if !okL || !okR { + continue + } else if shl.Op == bytecode.OP_SHR { + shl, shr = shr, shl + } + // + if shl.Op != bytecode.OP_SHL || shr.Op != bytecode.OP_SHR || shl.Left != shr.Left { + continue + } + // The rotated source must be stable across both shifts. + var ( + x = shl.Left + width = uint(or.Bitwidth) + ) + // + if writes[x] > 1 || regs[x].IsNative() || regs[x].Bitwidth().Unwrap() != width { + continue + } else if uint(shl.Bitwidth) != width || uint(shr.Bitwidth) != width { + continue + } + // + if rot, ok := matchRotation(shl, shr, width, regs, writes, defOf, consts); ok { + rots[or] = rot + dead[shl] = true + dead[shr] = true + } + } + } + // + return rots, dead +} + +// matchRotation classifies a (SHL, SHR) pair over source x and width w as a +// constant or dynamic rotation, per the rules described on scanRotations. +func matchRotation[W word.Word[W]](shl, shr *bytecode.Bitwise[W], width uint, + regs []descriptor.Register[W], + writes map[bytecode.RegisterId]uint, + defOf map[bytecode.RegisterId]Bytecode[W], + consts map[bytecode.RegisterId]W, +) (rotation, bool) { + // Shift amounts are always registers (constant operands are only + // supported for AND/OR/XOR). + shlAmt, shrAmt := shl.Right.AsRegister(), shr.Right.AsRegister() + // Constant case: both amounts known, summing to the width. + cl, okL := consts[shlAmt] + cr, okR := consts[shrAmt] + // + if okL && okR && cl.BitLen() <= 32 && cr.BitLen() <= 32 { + s, t := uint(cl.Uint64()), uint(cr.Uint64()) + // + if s+t == width { + return rotation{source: shl.Left, constAmount: s, isConst: true}, true + } + // + return rotation{}, false + } + // Dynamic case: power-of-two widths only, so that an amount register of + // width <= log2(width) always denotes an exact rotation. + if width == 0 || width&(width-1) != 0 { + return rotation{}, false + } + // + depth := uint(bits.Len(width - 1)) + // (x << n) | (x >> (w - n)): rotate left by n. + if n := shlAmt; isWidthMinus(shrAmt, width, n, writes, defOf, consts) && + isRotAmount(n, depth, regs, writes) { + return rotation{source: shl.Left, amount: n, op: bytecode.OP_SHL}, true + } + // (x << (w - n)) | (x >> n): rotate right by n. + if n := shrAmt; isWidthMinus(shlAmt, width, n, writes, defOf, consts) && + isRotAmount(n, depth, regs, writes) { + return rotation{source: shr.Left, amount: n, op: bytecode.OP_SHR}, true + } + // + return rotation{}, false +} + +// isRotAmount checks a dynamic rotation amount register: stable (at most one +// write), non-native, and narrow enough to enter the rotation chain directly +// (width in [1, depth]). +func isRotAmount[W word.Word[W]](n bytecode.RegisterId, depth uint, + regs []descriptor.Register[W], writes map[bytecode.RegisterId]uint, +) bool { + if writes[n] > 1 || regs[n].IsNative() { + return false + } + // + w := regs[n].Bitwidth().Unwrap() + // + return w >= 1 && w <= depth +} + +// isWidthMinus reports whether register m is structurally "width - n": a +// single-writer two-source SUB with zero constant whose first source holds +// the constant width and whose second source is n. +func isWidthMinus[W word.Word[W]](m bytecode.RegisterId, width uint, n bytecode.RegisterId, + writes map[bytecode.RegisterId]uint, + defOf map[bytecode.RegisterId]Bytecode[W], + consts map[bytecode.RegisterId]W, +) bool { + if writes[m] != 1 { + return false + } + // + sub, ok := defOf[m].(*bytecode.Arith[W]) + if !ok || sub.Op != bytecode.OP_SUB || len(sub.Target) != 1 || len(sub.Source) != 2 { + return false + } + // + if sub.Constant.Cmp64(0) != 0 || sub.Source[1] != n { + return false + } + // + w, ok := consts[sub.Source[0]] + // + return ok && w.Cmp64(uint64(width)) == 0 +} diff --git a/pkg/zkc/vm/internal/transform/split/concat.go b/pkg/zkc/vm/internal/transform/split/concat.go index 8b15b0f08..24c155d1e 100644 --- a/pkg/zkc/vm/internal/transform/split/concat.go +++ b/pkg/zkc/vm/internal/transform/split/concat.go @@ -66,21 +66,20 @@ type partCat struct { func Concat[W word.Word[W]](mapping descriptor.LimbsMap[W], alloc Allocator[W], insn *bytecode.Cat[W]) []Bytecode[W] { // Split into the initial set of chunks. - var chunks, context = initialiseConcatChunks(mapping, alloc, insn.Targets, insn.Sources) + var chunks, pre, context = initialiseConcatChunks(mapping, alloc, insn.Targets, insn.Sources) // Next, add carry lines as needed chunks = insertConcatCarryLines(mapping.Field(), alloc, chunks) // Convert chunks into assignments - return append(array.Map(chunks, concatAssignment[W]), context...) + return append(append(pre, array.Map(chunks, concatAssignment[W])...), context...) } // initialiseAddChunks splits the addition sources and constant into // least-significant-first chunks, then assigns target limbs to each chunk // according to the number of bits the corresponding RHS can produce. func initialiseConcatChunks[W word.Word[W]](mapping descriptor.LimbsMap[W], alloc Allocator[W], - targets, sources []RegisterId) ([]partCat, []Bytecode[W]) { + targets, sources []RegisterId) ([]partCat, []Bytecode[W], []Bytecode[W]) { // var ( - limbsMap = mapping.LimbsRegisterMap() // Split all source registers sourceLimbs = applyLimbsMapReversed(mapping, sources...) // Split all target registers @@ -91,18 +90,31 @@ func initialiseConcatChunks[W word.Word[W]](mapping descriptor.LimbsMap[W], allo sourceStack = RegisterStack[W]{sourceLimbs, alloc, nil} // codes []partCat + pre []Bytecode[W] ) // Keep going whilst we still have source registers for sourceStack.Size() > 0 { var ( rhs = sourceStack.SelectUpto(mapping.BandWidth()) - bitwidth = concatRhsBitwidth(mapping.Field(), rhs, limbsMap) + bitwidth = concatRhsBitwidth(mapping.Field(), rhs, alloc) // Keep target limbs intact. Splitting a target to match the RHS // exactly introduces a reconstruction bytecode which can sit after a // large control-flow join and make its path condition prohibitively // expensive to lower. lhs = targetStack.SelectUpto(bitwidth) ) + // A force-selected target limb can be wider than this chunk's sources: + // e.g. a u15 source limb whose successor would exceed the bandwidth, + // selected against a u16 target limb. Zero-padding mid-stream would + // shift every remaining source bit out of position, so instead borrow + // the missing low bits from the next source limb(s), destructing a limb + // when only part of it is needed. + if lhsWidth := descriptor.BitwidthOf(alloc, lhs...).Unwrap(); lhsWidth > bitwidth { + var borrows []Bytecode[W] + + rhs, borrows = borrowConcatSources(alloc, &sourceStack, rhs, bitwidth, lhsWidth) + pre = append(pre, borrows...) + } // allocate selected targets codes = append(codes, partCat{lhs, rhs}) } @@ -116,7 +128,46 @@ func initialiseConcatChunks[W word.Word[W]](mapping descriptor.LimbsMap[W], allo // Assert that we never create bytecodes for source registers. util.Assert(len(sourceStack.post) == 0, "internal failure") // - return codes, targetStack.post + return codes, pre, targetStack.post +} + +// borrowConcatSources widens a chunk's sources from bitwidth up to lhsWidth by +// consuming further limbs from the source stack: a limb no wider than the +// deficit is taken whole, while a wider one is destructed into a borrowed low +// part and a remainder (pushed back for the next chunk). It returns the +// widened sources and the destruct bytecodes (which must execute before the +// chunk assignments). If the stack empties first the deficit remains, which +// is fine: the chunk is then final and zero-extends. +func borrowConcatSources[W word.Word[W]](alloc Allocator[W], sourceStack *RegisterStack[W], + rhs []RegisterId, bitwidth, lhsWidth uint, +) ([]RegisterId, []Bytecode[W]) { + var pre []Bytecode[W] + // + for bitwidth < lhsWidth && sourceStack.Size() > 0 && + !descriptor.HasNativeRegisterId(sourceStack.stack, alloc) { + var ( + need = lhsWidth - bitwidth + next = sourceStack.Pop() + nextWidth = alloc.Register(next).Bitwidth().Unwrap() + ) + // + if nextWidth <= need { + rhs = append(rhs, next) + bitwidth += nextWidth + } else { + var ( + lo = alloc.Allocate("b", util.Some(need)) + hi = alloc.Allocate("b", util.Some(nextWidth-need)) + ) + // + pre = append(pre, bytecode.AddVec[W]([]RegisterId{lo, hi}, []RegisterId{next})) + rhs = append(rhs, lo) + sourceStack.stack = array.Prepend(hi, sourceStack.stack) + bitwidth += need + } + } + // + return rhs, pre } // insertConcatCarryLines allocates carry registers for chunks whose RHS diff --git a/testdata/zkc/unit/rotate_01.accepts b/testdata/zkc/unit/rotate_01.accepts new file mode 100644 index 000000000..abf7de614 --- /dev/null +++ b/testdata/zkc/unit/rotate_01.accepts @@ -0,0 +1,5 @@ +{ "data": "0x000100020008", "amount": "0b0001" } +{ "data": "0x800000010004", "amount": "0b0001" } +{ "data": "0xabcdabcd5e6d", "amount": "0b0000" } +{ "data": "0xabcdd5e65e6d", "amount": "0b1111" } +{ "data": "0x12341a0991a0", "amount": "0b0111" } diff --git a/testdata/zkc/unit/rotate_01.rejects b/testdata/zkc/unit/rotate_01.rejects new file mode 100644 index 000000000..61123bb15 --- /dev/null +++ b/testdata/zkc/unit/rotate_01.rejects @@ -0,0 +1,3 @@ +{ "data": "0x000100030008", "amount": "0b0001" } +{ "data": "0x000100020009", "amount": "0b0001" } +{ "data": "0xabcdd5e75e6d", "amount": "0b1111" } diff --git a/testdata/zkc/unit/rotate_01.zkc b/testdata/zkc/unit/rotate_01.zkc new file mode 100644 index 000000000..d1efd6d6d --- /dev/null +++ b/testdata/zkc/unit/rotate_01.zkc @@ -0,0 +1,23 @@ +pub input data(address:u2) -> (word:u16) +pub input amount(offset:u1) -> (n:u4) + +// prove that ((x << n) | (x >> (16 - n))) == y (i.e. rotate left), for both +// variable and constant rotation amounts +fn main() { + var x:u16 = data[0] + var y:u16 = data[1] + var z:u16 = data[2] + var n:u4 = amount[0] + var shift:u5 = 16 - (n as u5) + + // variable amount + if ((x << n) | (x >> shift)) != y { + fail + } + // constant amount + if ((x << 3) | (x >> 13)) != z { + fail + } + + return +}