JSON.cpp

Go to the documentation of this file.

//===-- JSON.cpp - llzk-witgen JSON conversion helpers ----------*- 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 "Errors.h"
#include "WitgenUtils.h"
#include "WitnessSelection.h"
 
#include "llzk/Dialect/Felt/IR/Types.h"
#include "llzk/Dialect/POD/IR/Attrs.h"
#include "llzk/Util/Compare.h"
 
#include <mlir/IR/Operation.h>
 
#include <llvm/ADT/SmallString.h>
#include <llvm/ADT/TypeSwitch.h>
#include <llvm/Support/FormatVariadic.h>
#include <llvm/Support/MathExtras.h>
#include <llvm/Support/raw_ostream.h>
 
using namespace mlir;
 
namespace llzk::witgen {
namespace {
 
static std::string renderJSON(const llvm::json::Value &value) {
  return llvm::formatv("{0:2}", value).str();
}
 
static llvm::StringRef jsonKind(const llvm::json::Value &value) {
  if (value.getAsNull()) {
    return "null";
  }
  if (value.getAsBoolean().has_value()) {
    return "bool";
  }
  if (value.getAsNumber().has_value()) {
    return "number";
  }
  if (value.getAsString()) {
    return "string";
  }
  if (value.getAsArray()) {
    return "array";
  }
  if (value.getAsObject()) {
    return "object";
  }
  return "unknown";
}
 
static std::string appendObjectPath(llvm::StringRef path, llvm::StringRef key) {
  llvm::SmallString<64> out(path);
  out += '.';
  out += key;
  return std::string(out);
}
 
static std::string appendIndexPath(llvm::StringRef path, size_t index) {
  llvm::SmallString<64> out(path);
  llvm::raw_svector_ostream os(out);
  os << '[' << index << ']';
  return std::string(out);
}
 
static void pushMismatch(
    llvm::SmallVectorImpl<JSONMismatch> &out, llvm::StringRef path, const llvm::Twine &message
) {
  out.push_back(JSONMismatch {path.str(), message.str()});
}
 
static void diffObjects(
    const llvm::json::Object &expected, const llvm::json::Object &actual,
    llvm::SmallVectorImpl<JSONMismatch> &out, llvm::StringRef path
) {
  for (const auto &kv : expected) {
    if (const llvm::json::Value *actualValue = actual.get(kv.first)) {
      diffJSON(kv.second, *actualValue, out, appendObjectPath(path, kv.first));
      continue;
    }
    pushMismatch(out, appendObjectPath(path, kv.first), "missing key");
  }
  for (const auto &kv : actual) {
    if (!expected.get(kv.first)) {
      pushMismatch(out, appendObjectPath(path, kv.first), "unexpected key");
    }
  }
}
 
static void diffArrays(
    const llvm::json::Array &expected, const llvm::json::Array &actual,
    llvm::SmallVectorImpl<JSONMismatch> &out, llvm::StringRef path
) {
  if (expected.size() != actual.size()) {
    pushMismatch(
        out, path,
        llvm::Twine("array length mismatch: expected ") + llvm::Twine(expected.size()) + ", got " +
            llvm::Twine(actual.size())
    );
  }
  size_t shared = std::min(expected.size(), actual.size());
  for (size_t i = 0; i < shared; ++i) {
    diffJSON(expected[i], actual[i], out, appendIndexPath(path, i));
  }
}
 
} // namespace

static llvm::Expected<int64_t> jsonToInt(const llvm::json::Value *json) {
  if (std::optional<int64_t> integer = json->getAsInteger()) {
    return *integer;
  }
  if (std::optional<llvm::StringRef> str = json->getAsString()) {
    int64_t value = 0;
    if (!str->getAsInteger(10, value)) {
      return value;
    }
  }
  return makeError("expected integer-compatible JSON value");
}

static llvm::Expected<llvm::DynamicAPInt>
jsonToFelt(const llvm::json::Value *json, const Field &field) {
  if (std::optional<llvm::StringRef> str = json->getAsString()) {
    auto value = parseDynamicAPInt(*str);
    if (!value) {
      return value.takeError();
    }
    return field.reduce(*value);
  }
  if (std::optional<int64_t> integer = json->getAsInteger()) {
    return field.reduce(*integer);
  }
  return makeError("expected felt value as JSON integer or decimal string");
}

static llvm::Expected<WitnessVal> parseJSONArray(
    const llvm::json::Value *json, array::ArrayType type, const Field &field, Operation *origin,
    size_t dimIndex = 0
) {
  const auto *jsonArray = json->getAsArray();
  if (!jsonArray) {
    return makeError("expected JSON array");
  }
 
  llvm::ArrayRef<int64_t> shape = type.getShape();
  if (dimIndex >= shape.size()) {
    return makeError("invalid array shape");
  }
  auto expectedSize = checkedShapeDimToSize(shape[dimIndex], "JSON array input");
  if (!expectedSize) {
    return expectedSize.takeError();
  }
  if (jsonArray->size() != *expectedSize) {
    return makeError("JSON array length does not match LLZK array dimension");
  }
 
  auto arrayValue = std::make_shared<ArrayValue>();
  arrayValue->type = type;
  if (dimIndex == shape.size() - 1) {
    arrayValue->elements.reserve(jsonArray->size());
    for (const llvm::json::Value &elem : *jsonArray) {
      auto parsed = parseJSONValue(&elem, type.getElementType(), field, origin);
      if (!parsed) {
        return parsed.takeError();
      }
      arrayValue->elements.push_back(*parsed);
    }
    return arrayValue;
  }
 
  arrayValue->elements.reserve(jsonArray->size());
  for (const llvm::json::Value &elem : *jsonArray) {
    auto parsed = parseJSONArray(&elem, type, field, origin, dimIndex + 1);
    if (!parsed) {
      return parsed.takeError();
    }
    auto subArray = asArray(*parsed);
    if (!subArray) {
      return subArray.takeError();
    }
    for (const WitnessVal &subElem : (*subArray)->elements) {
      arrayValue->elements.push_back(subElem);
    }
  }
  return arrayValue;
}
 
llvm::json::Value serializeFeltToJSON(const llvm::DynamicAPInt &value, const Field &field) {
  std::string rendered;
  llvm::raw_string_ostream(rendered) << field.reduce(value);
  return llvm::json::Value(rendered);
}

static llvm::Expected<llvm::json::Value> serializeJSONArray(
    const ArrayValueRef &arrayValue, array::ArrayType type, SymbolTableCollection &tables,
    Operation *origin, SerializationMode mode, size_t dimIndex = 0, size_t flatOffset = 0
) {
  llvm::json::Array jsonArray;
  llvm::ArrayRef<int64_t> shape = type.getShape();
  auto dimSize = checkedShapeDimToSize(shape[dimIndex], "JSON array output");
  if (!dimSize) {
    return dimSize.takeError();
  }
  if (dimIndex == shape.size() - 1) {
    for (size_t i = 0; i < *dimSize; ++i) {
      bool overflow = false;
      size_t elementOffset = llvm::SaturatingAdd(flatOffset, i, &overflow);
      if (overflow) {
        return makeError("JSON array output flat index would overflow size_t");
      }
      auto elem = serializeJSONValue(
          arrayValue->elements[elementOffset], type.getElementType(), tables, origin, mode
      );
      if (!elem) {
        return elem.takeError();
      }
      jsonArray.push_back(*elem);
    }
    return llvm::json::Value(std::move(jsonArray));
  }
 
  size_t subArraySize = 1;
  for (size_t i = dimIndex + 1; i < shape.size(); ++i) {
    auto nextDimSize = checkedShapeDimToSize(shape[i], "JSON array output");
    if (!nextDimSize) {
      return nextDimSize.takeError();
    }
    bool overflow = false;
    subArraySize = llvm::SaturatingMultiply(subArraySize, *nextDimSize, &overflow);
    if (overflow) {
      return makeError("JSON array output sub-array size would overflow size_t");
    }
  }
 
  for (size_t i = 0; i < *dimSize; ++i) {
    bool overflow = false;
    size_t nextOffset = llvm::SaturatingMultiplyAdd(i, subArraySize, flatOffset, &overflow);
    if (overflow) {
      return makeError("JSON array output flat offset would overflow size_t");
    }
    auto subArray =
        serializeJSONArray(arrayValue, type, tables, origin, mode, dimIndex + 1, nextOffset);
    if (!subArray) {
      return subArray.takeError();
    }
    jsonArray.push_back(*subArray);
  }
  return llvm::json::Value(std::move(jsonArray));
}

llvm::Expected<WitnessVal>
parseJSONValue(const llvm::json::Value *json, Type type, const Field &field, Operation *origin) {
  return llvm::TypeSwitch<Type, llvm::Expected<WitnessVal>>(type)
      .Case([&](felt::FeltType) -> llvm::Expected<WitnessVal> { return jsonToFelt(json, field); })
      .Case([&](array::ArrayType arrayType) -> llvm::Expected<WitnessVal> {
    return parseJSONArray(json, arrayType, field, origin);
  })
      .Case([&](pod::PodType) -> llvm::Expected<WitnessVal> {
    return makeError("pod JSON inputs are not supported in llzk-witgen v1");
  })
      .Case([&](component::StructType) -> llvm::Expected<WitnessVal> {
    return makeError("struct JSON inputs are not supported in llzk-witgen v1");
  })
      .Case([&](IndexType) -> llvm::Expected<WitnessVal> {
    auto integer = jsonToInt(json);
    if (!integer) {
      return integer.takeError();
    }
    return *integer;
  })
      .Case([&](IntegerType intType) -> llvm::Expected<WitnessVal> {
    if (intType.getWidth() == 1) {
      if (std::optional<bool> boolValue = json->getAsBoolean()) {
        return *boolValue;
      }
      auto integer = jsonToInt(json);
      if (!integer) {
        return integer.takeError();
      }
      return *integer != 0;
    }
    return makeError("only i1 integer JSON inputs are supported");
  }).Default([&](Type) -> llvm::Expected<WitnessVal> {
    return makeError("unsupported input type in llzk-witgen");
  });
}

llvm::Expected<llvm::json::Value> serializeJSONValue(
    const WitnessVal &value, Type type, SymbolTableCollection &tables, Operation *origin,
    SerializationMode mode
) {
  return llvm::TypeSwitch<Type, llvm::Expected<llvm::json::Value>>(type)
      .Case([&](felt::FeltType feltType) -> llvm::Expected<llvm::json::Value> {
    auto feltValue = asFelt(value);
    if (!feltValue) {
      return feltValue.takeError();
    }
    if (!feltType.hasField()) {
      return makeError("felt JSON output requires a declared field");
    }
    return serializeFeltToJSON(*feltValue, feltType.getField());
  })
      .Case([&](array::ArrayType arrayType) -> llvm::Expected<llvm::json::Value> {
    auto arrayValue = asArray(value);
    if (!arrayValue) {
      return arrayValue.takeError();
    }
    return serializeJSONArray(*arrayValue, arrayType, tables, origin, mode);
  })
      .Case([&](pod::PodType podType) -> llvm::Expected<llvm::json::Value> {
    auto podValue = asPod(value);
    if (!podValue) {
      return podValue.takeError();
    }
    llvm::json::Object result;
    for (pod::RecordAttr record : podType.getRecords()) {
      auto it = (*podValue)->records.find(record.getName().getValue());
      if (it == (*podValue)->records.end()) {
        return makeError("missing POD record during JSON serialization");
      }
      auto serialized = serializeJSONValue(it->second, record.getType(), tables, origin, mode);
      if (!serialized) {
        return serialized.takeError();
      }
      result[record.getName().getValue()] = *serialized;
    }
    return llvm::json::Value(std::move(result));
  })
      .Case([&](component::StructType structType) -> llvm::Expected<llvm::json::Value> {
    auto structValue = asStruct(value);
    if (!structValue) {
      return structValue.takeError();
    }
    auto defLookup = structType.getDefinition(tables, origin);
    if (failed(defLookup)) {
      return makeError("could not resolve struct type during JSON serialization");
    }
    llvm::json::Object result;
    for (component::MemberDefOp member : defLookup->get().getMemberDefs()) {
      auto it = (*structValue)->members.find(member.getSymName());
      if (it == (*structValue)->members.end()) {
        return makeError("missing struct member during JSON serialization");
      }
 
      if (mode == SerializationMode::PublicOutputsOnly) {
        if (!member.hasPublicAttr()) {
          continue;
        }
      } else {
        if (!memberIsSignal(defLookup->get(), member) &&
            !isa<component::StructType>(member.getType())) {
          continue;
        }
      }
 
      auto serialized = serializeJSONValue(it->second, member.getType(), tables, origin, mode);
      if (!serialized) {
        return serialized.takeError();
      }
      if (mode == SerializationMode::AllSignals && isa<component::StructType>(member.getType())) {
        auto *object = serialized->getAsObject();
        if (!object || object->empty()) {
          continue;
        }
      }
      result[member.getSymName()] = *serialized;
    }
    return llvm::json::Value(std::move(result));
  })
      .Case([&](IndexType) -> llvm::Expected<llvm::json::Value> {
    auto indexValue = asIndex(value);
    if (!indexValue) {
      return indexValue.takeError();
    }
    return llvm::json::Value(*indexValue);
  })
      .Case([&](IntegerType intType) -> llvm::Expected<llvm::json::Value> {
    if (intType.getWidth() != 1) {
      return makeError("only i1 integer JSON serialization is supported");
    }
    auto boolValue = asBool(value);
    if (!boolValue) {
      return boolValue.takeError();
    }
    return llvm::json::Value(*boolValue);
  }).Default([&](Type) -> llvm::Expected<llvm::json::Value> {
    return makeError("unsupported output type in llzk-witgen");
  });
}

llvm::Expected<llvm::json::Object> buildInputsJSONObject(
    ArrayRef<InputBinding> bindings, ArrayRef<WitnessVal> values, SymbolTableCollection &tables,
    Operation *origin
) {
  if (bindings.size() != values.size()) {
    return makeError("input binding count mismatch during witness JSON assembly");
  }
 
  llvm::json::Object result;
  for (auto [binding, value] : llvm::zip(bindings, values)) {
    auto serialized =
        serializeJSONValue(value, binding.type, tables, origin, SerializationMode::AllSignals);
    if (!serialized) {
      return serialized.takeError();
    }
    result[binding.name] = *serialized;
  }
  return result;
}

llvm::Expected<WitnessVal> extractValueAtPath(
    const WitnessVal &root, Type rootType, ArrayRef<std::string> path,
    SymbolTableCollection &tables, Operation *origin
) {
  if (path.empty()) {
    return root;
  }
 
  if (auto structType = dyn_cast<component::StructType>(rootType)) {
    auto structValue = asStruct(root);
    if (!structValue) {
      return structValue.takeError();
    }
    auto defLookup = structType.getDefinition(tables, origin);
    if (failed(defLookup)) {
      return makeError("could not resolve struct type while extracting witness value");
    }
    for (component::MemberDefOp member : defLookup->get().getMemberDefs()) {
      if (member.getSymName() != path.front()) {
        continue;
      }
      auto it = (*structValue)->members.find(member.getSymName());
      if (it == (*structValue)->members.end()) {
        return makeError("missing struct member while extracting witness value");
      }
      return extractValueAtPath(it->second, member.getType(), path.drop_front(), tables, origin);
    }
    return makeError("unknown struct member while extracting witness value");
  }
 
  if (auto podType = dyn_cast<pod::PodType>(rootType)) {
    auto podValue = asPod(root);
    if (!podValue) {
      return podValue.takeError();
    }
    for (pod::RecordAttr record : podType.getRecords()) {
      if (record.getName().getValue() != path.front()) {
        continue;
      }
      auto it = (*podValue)->records.find(record.getName().getValue());
      if (it == (*podValue)->records.end()) {
        return makeError("missing POD record while extracting witness value");
      }
      return extractValueAtPath(it->second, record.getType(), path.drop_front(), tables, origin);
    }
    return makeError("unknown POD record while extracting witness value");
  }
 
  return makeError("extra witness path components for non-aggregate value");
}
 
void diffJSON(
    const llvm::json::Value &expected, const llvm::json::Value &actual,
    llvm::SmallVectorImpl<JSONMismatch> &out, llvm::StringRef path
) {
  if (expected.kind() != actual.kind()) {
    pushMismatch(
        out, path,
        llvm::Twine("type mismatch: expected ") + jsonKind(expected) + ", got " + jsonKind(actual)
    );
    return;
  }
 
  if (const auto *expectedObject = expected.getAsObject()) {
    diffObjects(*expectedObject, *actual.getAsObject(), out, path);
    return;
  }
  if (const auto *expectedArray = expected.getAsArray()) {
    diffArrays(*expectedArray, *actual.getAsArray(), out, path);
    return;
  }
  if (expected == actual) {
    return;
  }
 
  pushMismatch(
      out, path,
      llvm::Twine("value mismatch: expected ") + renderJSON(expected) + ", got " +
          renderJSON(actual)
  );
}
 
void printJSONMismatches(llvm::raw_ostream &os, llvm::ArrayRef<JSONMismatch> mismatches) {
  for (const JSONMismatch &mismatch : mismatches) {
    os << mismatch.path << ": " << mismatch.message << '\n';
  }
}
 
} // namespace llzk::witgen