Use cycle counter in Keccak circuit for provable determinism

Author: jacobdweightman

zirgen/circuit/keccak/cycle_counter.zir

@@ -0,0 +1,34 @@
+// RUN: true
+
+import is_zero;
+
+extern GetCycle(): Val;
+
+// RISC Zero STARKs are equivalent up to rotation of the trace. This component
+// counts cycles from zero up to the power of two, with constraints that
+// guarantee a unique "cycle zero."
+
+// It uses a global that stores the total number of cycles, which the verifier
+// should check matches the intended trace length. If the global doesn't match,
+// the verifier should reject, so assume that it matches. Then we either have a
+// cycle zero or we don't; if we do, then the next cycle must be 1, then 2, and
+// so on because of the constraint cycle = cycle@1 + 1. After total_cycles, we
+// have gone over the whole trace, which means the cycle before cycle 0 must be
+// total_cycles - 1. If we don't, then we always have that cycle = cycle@1. But
+// because the trace is cyclic, the cycle number must go down on some cycle, so
+// a constraint must have been violated.
+component CycleCounter() {
+  global total_cycles : NondetReg;
+
+  cycle := NondetReg(GetCycle());
+  public is_first_cycle := IsZero(cycle);
+
+  [is_first_cycle, 1-is_first_cycle] -> ({
+    // First cycle; previous cycle should be the last cycle.
+    cycle@1 = total_cycles - 1;
+  }, {
+    // Not first cycle; cycle number should advance by one for every row.
+    cycle = cycle@1 + 1;
+  });
+  cycle
+}

zirgen/circuit/keccak/predicates.cpp

@@ -27,13 +27,14 @@ template <typename Func>
 void addZirgenLift(Module& module, const std::string name, const std::string path, Func func) {
   auto circuit = getInterfaceZirgen(module.getModule().getContext(), path);
   for (size_t po2 = 14; po2 < 19; ++po2) {
+    size_t totalCycles = 1 << po2;
     module.addFunc<3>(name + "_" + std::to_string(po2),
                       {gbuf(recursion::kOutSize), ioparg(), ioparg()},
                       [&](Buffer out, ReadIopVal rootIop, ReadIopVal zirgenSeal) {
                         DigestVal root = rootIop.readDigests(1)[0];
                         VerifyInfo info = zirgen::verify::verify(zirgenSeal, po2, *circuit);
                         llvm::ArrayRef inStream(info.out);
-                        DigestVal outData = func(inStream);
+                        DigestVal outData = func(inStream, totalCycles);
                         std::vector<Val> outStream;
                         writeSha(outData, outStream);
                         doExtern("write", "", 0, outStream);
@@ -57,9 +58,15 @@ int main(int argc, char* argv[]) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "keccak predicates");
 
   Module module;
-  addZirgenLift(module, "keccak_lift", keccakIR.getValue(), [](llvm::ArrayRef<Val>& inStream) {
-    return readSha(inStream);
-  });
+  addZirgenLift(module,
+                "keccak_lift",
+                keccakIR.getValue(),
+                [](llvm::ArrayRef<Val>& inStream, size_t expectedTotalCycles) {
+                  auto sha = readSha(inStream);
+                  Val totalCycles = readVal(inStream);
+                  eq(totalCycles, expectedTotalCycles);
+                  return sha;
+                });
 
   module.optimize();
   module.getModule().walk([&](mlir::func::FuncOp func) { zirgen::emitRecursion(outputDir, func); });

zirgen/circuit/keccak/top.zir

@@ -1,6 +1,6 @@
 // RUN: true
 
-import is_zero;
+import cycle_counter;
 import keccak;
 import sha2;
 
@@ -276,7 +276,6 @@ component ShaNextBlockCycle(back1: TopState) {
   topState
 }
 
-extern IsFirstCycle(): Val;
 extern GetPreimage(idx: Val): Val;
 extern NextPreimage(): Val;
 
@@ -474,10 +473,10 @@ component WrapOneHot(oneHot: OneHot<12>) {
 component Top() {
   global finalDigest: DigestReg;
 
-  isFirst := NondetReg(IsFirstCycle());
+  cycle := CycleCounter();
   cycleMux : WrapOneHot;
   controlState : ControlState;
-  controlState := if (isFirst) {
+  controlState := if (cycle.is_first_cycle) {
     [email protected] = CycleTypeShutdown();
     ControlState(CycleTypeInit(), 0, 0, 0)
   } else {

zirgen/circuit/predicates/predicates.cpp

@@ -25,7 +25,7 @@ using Zll::DigestKind;
 
 constexpr size_t kMaxInsnCycles = 2000; // TODO(flaub): update this with precise value.
 
-static Val readVal(llvm::ArrayRef<Val>& stream) {
+Val readVal(llvm::ArrayRef<Val>& stream) {
   assert(stream.size() >= 1);
   Val out = stream[0];
   stream = stream.drop_front();

zirgen/circuit/predicates/predicates.h

@@ -118,6 +118,7 @@ ReceiptClaim join(ReceiptClaim in1, ReceiptClaim in2);
 ReceiptClaim identity(ReceiptClaim in);
 ReceiptClaim resolve(ReceiptClaim cond, Assumption assum, DigestVal tail, DigestVal journal);
 
+Val readVal(llvm::ArrayRef<Val>& stream);
 DigestVal readSha(llvm::ArrayRef<Val>& stream, bool longDigest = false);
 void writeSha(DigestVal val, std::vector<Val>& stream);