llzk-witgen.cpp

Go to the documentation of this file.

//===-- llzk-witgen.cpp - LLZK witness generation tool ----------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#include "JSON.h"
#include "WitgenDriver.h"
#include "Wtns.h"
#include "tools/config.h"
 
#include "llzk/Dialect/Array/IR/Dialect.h"
#include "llzk/Dialect/Bool/IR/Dialect.h"
#include "llzk/Dialect/Cast/IR/Dialect.h"
#include "llzk/Dialect/Constrain/IR/Dialect.h"
#include "llzk/Dialect/DialectRegistration.h"
#include "llzk/Dialect/Felt/IR/Dialect.h"
#include "llzk/Dialect/Function/IR/Dialect.h"
#include "llzk/Dialect/Global/IR/Dialect.h"
#include "llzk/Dialect/Include/IR/Dialect.h"
#include "llzk/Dialect/Include/Util/IncludeHelper.h"
#include "llzk/Dialect/LLZK/IR/Dialect.h"
#include "llzk/Dialect/POD/IR/Dialect.h"
#include "llzk/Dialect/Polymorphic/IR/Dialect.h"
#include "llzk/Dialect/RAM/IR/Dialect.h"
#include "llzk/Dialect/String/IR/Dialect.h"
#include "llzk/Dialect/Struct/IR/Dialect.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
 
#include <mlir/Dialect/Arith/IR/Arith.h>
#include <mlir/Dialect/ControlFlow/IR/ControlFlowOps.h>
#include <mlir/Dialect/Func/Extensions/InlinerExtension.h>
#include <mlir/Dialect/Func/IR/FuncOps.h>
#include <mlir/Dialect/MemRef/IR/MemRef.h>
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/Dialect/SMT/IR/SMTDialect.h>
#include <mlir/IR/BuiltinDialect.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/Diagnostics.h>
#include <mlir/IR/DialectRegistry.h>
#include <mlir/Parser/Parser.h>
 
#include <llvm/ADT/StringExtras.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/Support/JSON.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/PrettyStackTrace.h>
#include <llvm/Support/Signals.h>
#include <llvm/Support/SourceMgr.h>
 
#include <cstdlib>
#include <exception>
 
using namespace mlir;

static int runMain(int argc, char **argv) {
  llvm::cl::opt<std::string> InputFilename(llvm::cl::Positional, llvm::cl::Required);
  llvm::cl::opt<std::string> InputsFilename(
      "inputs", llvm::cl::Required, llvm::cl::desc("JSON input file")
  );
  llvm::cl::list<std::string> IncludeDirs(
      "I", llvm::cl::desc("Directory of include files"), llvm::cl::value_desc("directory"),
      llvm::cl::Prefix
  );
  llvm::cl::opt<std::string> BackendName(
      "backend", llvm::cl::desc("Execution backend: interpreter or execution-engine"),
      llvm::cl::init("interpreter")
  );
  llvm::cl::opt<std::string> OutputScopeName(
      "output-scope", llvm::cl::desc("Output scope: public or full-witness"),
      llvm::cl::init("public")
  );
  llvm::cl::opt<std::string> UninitializedBehaviorName(
      "uninitialized-behavior",
      llvm::cl::desc("Uninitialized value behavior: zero, random, or fail"), llvm::cl::init("zero")
  );
  llvm::cl::opt<uint64_t> UninitializedSeed(
      "uninitialized-seed", llvm::cl::desc("Seed for random uninitialized values")
  );
  llvm::cl::opt<bool> DumpJITCore("dump-jit-core", llvm::cl::desc("Print the pre-LLVM JIT module"));
  llvm::cl::opt<bool> DumpJITLLVM(
      "dump-jit-llvm", llvm::cl::desc("Print the post-LLVM JIT module")
  );
  llvm::cl::opt<std::string> CheckOutputFilename(
      "check-output", llvm::cl::desc("JSON file with expected witgen output")
  );
  llvm::cl::opt<std::string> WtnsOutputFilename(
      "output-wtns", llvm::cl::desc("Write a snarkjs-compatible .wtns file")
  );
 
  llvm::sys::PrintStackTraceOnErrorSignal(llvm::StringRef());
  llvm::setBugReportMsg(
      "PLEASE submit a bug report to " BUG_REPORT_URL
      " and include the crash backtrace, relevant LLZK files, and associated run script(s).\n"
  );
 
  llvm::cl::ParseCommandLineOptions(
      argc, argv,
      "llzk-witgen: execute LLZK compute semantics and emit JSON public outputs.\n"
      "Note: llzk-witgen v1 ignores constrain() and traps on bool.assert.\n"
  );
 
  DialectRegistry registry;
  llzk::registerDialects(registry);
  llzk::registerInliningExtensions(registry);
  mlir::func::registerInlinerExtension(registry);
  registry.insert<
      mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect, mlir::func::FuncDialect,
      mlir::memref::MemRefDialect, mlir::scf::SCFDialect>();
 
  MLIRContext context;
  context.appendDialectRegistry(registry);
  context.loadAllAvailableDialects();
 
  if (failed(llzk::GlobalSourceMgr::get().setup(IncludeDirs))) {
    return EXIT_FAILURE;
  }
 
  auto sourceBuffer = llvm::MemoryBuffer::getFileOrSTDIN(InputFilename);
  if (!sourceBuffer) {
    llvm::errs() << sourceBuffer.getError().message() << '\n';
    return EXIT_FAILURE;
  }
 
  ParserConfig parserConfig(&context);
  llvm::SourceMgr sourceMgr;
  sourceMgr.AddNewSourceBuffer(std::move(*sourceBuffer), llvm::SMLoc());
  // Keep source diagnostics active through preprocessing and witness generation.
  SourceMgrDiagnosticHandler diagnosticHandler(sourceMgr, &context);
  OwningOpRef<ModuleOp> moduleOp = parseSourceFile<ModuleOp>(sourceMgr, parserConfig);
  if (!moduleOp) {
    return EXIT_FAILURE;
  }
 
  auto buffer = llvm::MemoryBuffer::getFileOrSTDIN(InputsFilename);
  if (!buffer) {
    llvm::errs() << buffer.getError().message() << '\n';
    return EXIT_FAILURE;
  }
 
  auto parsed = llvm::json::parse(buffer.get()->getBuffer());
  if (!parsed) {
    llvm::errs() << "failed to parse JSON input: " << llvm::toString(parsed.takeError()) << '\n';
    return EXIT_FAILURE;
  }
 
  llzk::witgen::WitgenOptions options;
  if (BackendName == "execution-engine") {
    options.backend = llzk::witgen::Backend::ExecutionEngine;
  } else if (BackendName == "interpreter") {
    options.backend = llzk::witgen::Backend::Interpreter;
  } else {
    llvm::errs() << "unknown backend: " << BackendName << '\n';
    return EXIT_FAILURE;
  }
  if (OutputScopeName == "full-witness") {
    options.outputScope = llzk::witgen::OutputScope::FullWitness;
  } else if (OutputScopeName == "public") {
    options.outputScope = llzk::witgen::OutputScope::Public;
  } else {
    llvm::errs() << "unknown output scope: " << OutputScopeName << '\n';
    return EXIT_FAILURE;
  }
  if (UninitializedBehaviorName == "zero") {
    options.uninitializedBehavior = llzk::witgen::UninitializedBehavior::Zero;
  } else if (UninitializedBehaviorName == "random") {
    options.uninitializedBehavior = llzk::witgen::UninitializedBehavior::Random;
  } else if (UninitializedBehaviorName == "fail") {
    options.uninitializedBehavior = llzk::witgen::UninitializedBehavior::Fail;
  } else {
    llvm::errs() << "unknown uninitialized behavior: " << UninitializedBehaviorName << '\n';
    return EXIT_FAILURE;
  }
  if (UninitializedSeed.getNumOccurrences() > 0) {
    options.randomSeed = UninitializedSeed;
  }
  options.inlineIncludes = true;
  options.dumpJITCore = DumpJITCore;
  options.dumpJITLLVM = DumpJITLLVM;
  if (WtnsOutputFilename.getNumOccurrences() > 0) {
    if (WtnsOutputFilename == "-") {
      llvm::errs() << "--output-wtns does not support stdout; specify an output file\n";
      return EXIT_FAILURE;
    }
    if (OutputScopeName.getNumOccurrences() > 0 && OutputScopeName != "full-witness") {
      llvm::errs() << "--output-wtns conflicts with --output-scope=" << OutputScopeName
                   << "; use --output-scope=full-witness\n";
      return EXIT_FAILURE;
    }
    // WTNS contains every R1CS wire, including ordinary private felt members.
    options.outputScope = llzk::witgen::OutputScope::R1CSWitness;
  }
 
  auto result = llzk::witgen::runWitgen(*moduleOp, *parsed, options);
  if (!result) {
    llvm::errs() << "llzk-witgen error: " << llvm::toString(result.takeError()) << '\n';
    return EXIT_FAILURE;
  }
 
  if (WtnsOutputFilename.getNumOccurrences() > 0) {
    // The WTNS header records the field modulus and element width. Re-read the
    // unique field from the preprocessed module so the header and generated
    // values necessarily describe the same field.
    llzk::FieldSet fields;
    if (failed(llzk::collectFields(moduleOp->getOperation(), fields)) || fields.size() != 1) {
      llvm::errs() << "llzk-witgen error: .wtns output requires exactly one field\n";
      return EXIT_FAILURE;
    }
    const llzk::Field &field = (*fields.begin()).get();
    if (auto error = llzk::witgen::writeWtns(*moduleOp, *result, field, WtnsOutputFilename)) {
      llvm::errs() << "llzk-witgen error: " << llvm::toString(std::move(error)) << '\n';
      return EXIT_FAILURE;
    }
  }
 
  if (CheckOutputFilename.getNumOccurrences() > 0) {
    auto expectedBuffer = llvm::MemoryBuffer::getFileOrSTDIN(CheckOutputFilename);
    if (!expectedBuffer) {
      llvm::errs() << expectedBuffer.getError().message() << '\n';
      return EXIT_FAILURE;
    }
 
    auto expected = llvm::json::parse(expectedBuffer.get()->getBuffer());
    if (!expected) {
      llvm::errs() << "failed to parse expected JSON output: "
                   << llvm::toString(expected.takeError()) << '\n';
      return EXIT_FAILURE;
    }
 
    llvm::SmallVector<llzk::witgen::JSONMismatch> mismatches;
    llzk::witgen::diffJSON(*expected, *result, mismatches);
    if (!mismatches.empty()) {
      llvm::errs() << "llzk-witgen output mismatch:\n";
      llzk::witgen::printJSONMismatches(llvm::errs(), mismatches);
      return EXIT_FAILURE;
    }
 
    llvm::outs() << "output matched expected JSON\n";
    return EXIT_SUCCESS;
  }
 
  llvm::outs() << llvm::formatv("{0:2}", *result) << '\n';
  return EXIT_SUCCESS;
}
 
int main(int argc, char **argv) noexcept {
  try {
    return runMain(argc, argv);
  } catch (const std::exception &ex) {
    llvm::errs() << "llzk-witgen: unhandled exception: " << ex.what() << '\n';
  } catch (...) {
    llvm::errs() << "llzk-witgen: unhandled non-standard exception\n";
  }
  return EXIT_FAILURE;
}