22#include <mlir/IR/AsmState.h>
38compareDynamicAPInt(
const llvm::DynamicAPInt &lhs,
const llvm::DynamicAPInt &rhs) {
40 return std::strong_ordering::less;
43 return std::strong_ordering::greater;
45 return std::strong_ordering::equal;
48std::strong_ordering compareStringRef(llvm::StringRef lhs, llvm::StringRef rhs) {
49 int cmp = lhs.compare(rhs);
51 return std::strong_ordering::less;
54 return std::strong_ordering::greater;
56 return std::strong_ordering::equal;
62LogicalResult printSourceStylePath(raw_ostream &os, llvm::ArrayRef<SourceRefIndex> path) {
63 for (
const auto &idx : path) {
65 os <<
'.' << idx.getMember().getName();
68 if (idx.isPodRecord()) {
69 os <<
'.' << idx.getPodRecordName();
72 if (idx.isIndex() || idx.isIndexRange()) {
73 os <<
'[' << idx <<
']';
82compareSourceRefPaths(llvm::ArrayRef<SourceRefIndex> lhs, llvm::ArrayRef<SourceRefIndex> rhs) {
83 for (
size_t i = 0; i < lhs.size() && i < rhs.size(); i++) {
84 if (
auto cmp = lhs[i] <=> rhs[i];
cmp != std::strong_ordering::equal) {
88 return lhs.size() <=> rhs.size();
104 if (ShapedType::isDynamic(int64_t(high))) {
107 os << low <<
':' << high;
115 return std::strong_ordering::less;
118 return std::strong_ordering::greater;
120 return std::strong_ordering::equal;
131 if (
auto cmp = compareDynamicAPInt(ll, rl);
cmp != std::strong_ordering::equal) {
134 return compareDynamicAPInt(lu, ru);
138 return std::strong_ordering::less;
141 return std::strong_ordering::greater;
144 return std::strong_ordering::less;
147 return std::strong_ordering::greater;
150 return std::strong_ordering::less;
152 return std::strong_ordering::greater;
162 unsigned requiredBits = idx.getSignificantBits();
163 auto hash = llvm::hash_value(idx.trunc(requiredBits));
167 return llvm::hash_value(std::get<0>(r)) ^ llvm::hash_value(std::get<1>(r));
186 return lhsLow < rhsHigh && rhsLow < lhsHigh;
193SourceRef::SortCategory SourceRef::getSortCategory()
const {
194 if (isBlockArgument()) {
195 return SortCategory::BlockArgument;
197 if (isCreateStructOp()) {
198 return SortCategory::CreateStruct;
201 return SortCategory::NonDet;
204 return SortCategory::RootResult;
206 if (isTemplateConstant()) {
207 return SortCategory::TemplateConstant;
209 if (isConstantIndex()) {
210 return SortCategory::ConstantIndex;
212 if (isConstantFelt()) {
213 return SortCategory::ConstantFelt;
216 llvm::errs() << *
this <<
'\n';
217 llvm_unreachable(
"unhandled SourceRef sort category");
220StringRef SourceRef::getTemplateConstantName()
const {
221 auto constantVal = getConstant();
222 ensure(succeeded(constantVal),
"template constant must be constant");
223 auto constRead = llvm::dyn_cast<ConstReadOp>(constantVal->getDefiningOp());
224 ensure(constRead,
"template constant must be backed by const.read");
225 return constRead.getConstName();
229SourceRef::compareWithinCategory(
const SourceRef &rhs, SortCategory category)
const {
231 case SortCategory::BlockArgument: {
232 if (
auto cmp = *getInputNum() <=> *rhs.getInputNum();
cmp != std::strong_ordering::equal) {
235 if (
auto cmp = getAsOpaquePointer() <=> rhs.getAsOpaquePointer();
236 cmp != std::strong_ordering::equal) {
239 return compareSourceRefPaths(getPath(), rhs.getPath());
241 case SortCategory::CreateStruct:
242 case SortCategory::NonDet:
243 case SortCategory::RootResult: {
244 if (
auto cmp = getAsOpaquePointer() <=> rhs.getAsOpaquePointer();
245 cmp != std::strong_ordering::equal) {
248 return compareSourceRefPaths(getPath(), rhs.getPath());
250 case SortCategory::TemplateConstant: {
251 if (
auto cmp = compareStringRef(getTemplateConstantName(), rhs.getTemplateConstantName());
252 cmp != std::strong_ordering::equal) {
255 return getAsOpaquePointer() <=> rhs.getAsOpaquePointer();
257 case SortCategory::ConstantIndex:
258 return compareDynamicAPInt(*getConstantIndexValue(), *rhs.getConstantIndexValue());
259 case SortCategory::ConstantFelt:
260 return compareDynamicAPInt(*getConstantFeltValue(), *rhs.getConstantFeltValue());
263 llvm_unreachable(
"unhandled SourceRef category compare");
273SymbolLookupResult<StructDefOp>
281 return std::move(*sDef);
284std::vector<SourceRef>
286 std::vector<SourceRef> res = {root};
287 for (
const SourceRef &child : root.getAllChildren(tables,
mod)) {
288 auto recursiveChildren = getAllSourceRefs(tables,
mod, child);
289 res.insert(res.end(), recursiveChildren.begin(), recursiveChildren.end());
295 std::vector<SourceRef> res;
298 structDef == fnOp->getParentOfType<
StructDefOp>(),
"function must be within the given struct"
302 ensure(succeeded(modOp),
"could not lookup module from struct " + Twine(structDef.getName()));
304 SymbolTableCollection tables;
305 for (
auto a : fnOp.getArguments()) {
307 res.insert(res.end(), argRes.begin(), argRes.end());
314 auto createOp = dyn_cast_if_present<CreateStructOp>(selfVal.getDefiningOp());
315 ensure(createOp,
"self value should originate from struct.new operation");
316 auto selfRes =
getAllSourceRefs(tables, modOp.value(), SourceRef(createOp));
317 res.insert(res.end(), selfRes.begin(), selfRes.end());
324 std::vector<SourceRef> res;
327 memberDef->getParentOfType<
StructDefOp>() == structDef,
328 "Member " + Twine(memberDef.getName()) +
" is not a member of struct " +
329 Twine(structDef.getName())
332 ensure(succeeded(modOp),
"could not lookup module from struct " + Twine(structDef.getName()));
335 BlockArgument self = constrainFnOp.getArguments().front();
336 SourceRef memberRef = SourceRef(self, {
SourceRefIndex(memberDef)});
338 SymbolTableCollection tables;
343 Type currTy = value.getType();
344 for (
const auto &idx :
getPath()) {
345 if (idx.isMember()) {
346 currTy = idx.getMember().getType();
349 if (idx.isPodRecord()) {
350 auto podTy = dyn_cast<pod::PodType>(currTy);
351 ensure(
static_cast<bool>(podTy),
"SourceRef pod record requires a pod-typed base");
352 auto lookup = podTy.getRecord(idx.getPodRecordName(), [ctx = value.getContext()]() {
353 return mlir::emitError(
354 mlir::UnknownLoc::get(ctx),
"SourceRef references a missing pod record"
357 ensure(succeeded(lookup),
"SourceRef references a missing pod record");
362 auto arrTy = dyn_cast<ArrayType>(currTy);
363 ensure(
static_cast<bool>(arrTy),
"SourceRef array index requires an array-typed base");
364 currTy = arrTy.getSelectionType(1);
375 auto prefixPath = prefix.
getPath();
376 if (value != prefix.value || pathRef.size() < prefixPath.size()) {
379 for (
size_t i = 0; i < prefixPath.size(); i++) {
380 if (pathRef[i] != prefixPath[i]) {
388 auto getSelfStruct = [](
const SourceRef &ref) ->
StructDefOp {
389 if (
auto createOp = dyn_cast_if_present<CreateStructOp>(ref.value.getDefiningOp())) {
390 auto func = createOp->getParentOfType<
FuncDefOp>();
391 if (!func || !func.isStructCompute() || func.getSelfValueFromCompute() != ref.value) {
396 auto blockArg = ref.getBlockArgument();
397 if (failed(blockArg)) {
400 auto func = dyn_cast_if_present<FuncDefOp>(blockArg->getOwner()->getParentOp());
401 return func && func.isStructConstrain() && func.getSelfValueFromConstrain() == *blockArg
405 bool sameRoot = value == rhs.value;
409 sameRoot = lhsStruct && lhsStruct == rhsStruct;
414 return llvm::all_of(llvm::zip(path, rhs.path), [](
const auto &indices) {
415 return std::get<0>(indices).overlaps(std::get<1>(indices));
420 llvm::SmallVector<SourceRefIndex> selections;
422 return index.isIndex() || index.isIndexRange();
425 SourceRef result = *
this;
426 size_t dimension = 0;
431 if (dimension < selections.size() && index.
isIndexRange() && selections[dimension].isIndex() &&
432 index.
overlaps(selections[dimension])) {
433 index = selections[dimension];
446 auto prefixPath = prefix.
getPath();
447 suffix.reserve(pathRef.size() - prefixPath.size());
448 for (
size_t i = prefixPath.size(); i < pathRef.size(); i++) {
449 suffix.push_back(pathRef[i]);
459 if (failed(suffix)) {
463 SourceRef newSignalUsage = other;
466 pathRef.insert(pathRef.end(), suffix->begin(), suffix->end());
469 return newSignalUsage;
475 std::vector<SourceRef> res;
477 for (int64_t i = 0; i < arrayTy.getDimSize(0); i++) {
479 ensure(succeeded(childRef),
"array children require a rooted SourceRef");
480 res.push_back(*childRef);
490 std::vector<SourceRef> res;
498 auto structDefCopy = structDefRes;
500 tables, SymbolRefAttr::get(f.getContext(), f.getSymNameAttr()), std::move(structDefCopy),
503 ensure(succeeded(memberLookup),
"could not get SymbolLookupResult of existing MemberDefOp");
505 ensure(succeeded(childRef),
"struct children require a rooted SourceRef");
508 res.push_back(*childRef);
514 std::vector<SourceRef> res;
517 ensure(succeeded(childRef),
"pod children require a rooted SourceRef");
518 res.push_back(*childRef);
523std::vector<SourceRef>
526 if (
auto structTy = dyn_cast<StructType>(ty)) {
528 }
else if (
auto podTy = dyn_cast<pod::PodType>(ty)) {
530 }
else if (
auto arrayType = dyn_cast<ArrayType>(ty)) {
537static void printCallResultFallback(raw_ostream &os,
function::CallOp callOp, Value value) {
540 Operation *printScope = callOp.getOperation();
541 if (
auto funcOp = callOp->getParentOfType<
FuncDefOp>()) {
542 printScope = funcOp.getOperation();
546 AsmState state(printScope);
547 value.printAsOperand(os, state);
551static bool shouldPrintNamedCallResult(
552 function::CallOp callOp, OpResult callResult, function::FuncDefOp calleeFunc
554 auto resName = calleeFunc.getResNameAttr(callResult.getResultNumber());
559 auto parentFunc = callOp->getParentOfType<FuncDefOp>();
564 bool foundThisCall =
false;
565 bool foundDuplicate =
false;
566 parentFunc.walk([&](function::CallOp otherCall) {
567 if (foundDuplicate) {
568 return WalkResult::interrupt();
571 auto otherFunc = llvm::dyn_cast_if_present<FuncDefOp>(otherCall.resolveCallable());
573 return WalkResult::advance();
575 for (Value otherValue : otherCall->getResults()) {
576 auto otherResult = llvm::cast<OpResult>(otherValue);
577 auto otherResName = otherFunc.getResNameAttr(otherResult.getResultNumber());
578 if (!otherResName || otherResName->getValue() != resName->getValue()) {
581 if (otherResult == callResult) {
582 foundThisCall =
true;
585 foundDuplicate =
true;
586 return WalkResult::interrupt();
588 return WalkResult::advance();
591 return foundThisCall && !foundDuplicate;
600 auto constRead = getDefiningOp<ConstReadOp>();
601 ensure(succeeded(constRead),
"template constant should be backed by a const.read op");
602 auto structDefOp = (*constRead)->getParentOfType<
StructDefOp>();
603 ensure(structDefOp,
"struct template should have a struct parent");
604 os <<
'@' << structDefOp.getName() <<
"<[@" << constRead->getConstName() <<
"]>";
610 auto funcOp = llvm::dyn_cast<FuncDefOp>(blockArg.getOwner()->getParentOp());
613 if (funcOp && funcOp.isStructConstrain() && funcOp.getSelfValueFromConstrain() == blockArg) {
616 std::optional<StringAttr> argName;
618 argName = funcOp.getArgNameAttr(blockArg.getArgNumber());
621 os << argName->getValue();
630 auto callResult = llvm::cast<OpResult>(value);
632 if (succeeded(callee)) {
634 if (calleeFunc && shouldPrintNamedCallResult(callOp, callResult, calleeFunc)) {
635 auto resName = *calleeFunc.
getResNameAttr(callResult.getResultNumber());
636 os << resName.getValue();
638 printCallResultFallback(os, callOp, value);
641 printCallResultFallback(os, callOp, value);
645 OpPrintingFlags flags;
646 value.printAsOperand(os, flags);
649 auto res = printSourceStylePath(os,
getPath());
650 ensure(succeeded(res),
"unhandled path print case");
660 return constant == rhs.constant && value == rhs.value && llvm::equal(
getPath(), rhs.
getPath());
665 auto lhsCategory = getSortCategory();
666 auto rhsCategory = rhs.getSortCategory();
667 if (
auto cmp = lhsCategory <=> rhsCategory;
cmp != std::strong_ordering::equal) {
670 return compareWithinCategory(rhs, lhsCategory);
677 return llvm::hash_value(val.getAsOpaquePointer());
681 "unhandled SourceRef hash case"
684 size_t hash = llvm::hash_value(val.getAsOpaquePointer());
685 for (
const auto &f : val.
getPath()) {
686 hash = llvm::hash_combine(hash, f.getHash());
700 insert(rhs.begin(), rhs.end());
706 std::vector<SourceRef> sortedRefs(rhs.begin(), rhs.end());
707 std::sort(sortedRefs.begin(), sortedRefs.end());
708 for (
auto it = sortedRefs.begin(); it != sortedRefs.end();) {
711 if (it != sortedRefs.end()) {
This file implements helper methods for constructing DynamicAPInts.
Shared utility function implementations for LLZK lowering passes.
This file defines methods symbol lookup across LLZK operations and included files.
Defines an index into an LLZK object.
std::strong_ordering operator<=>(const SourceRefIndex &rhs) const
bool isIndexRange() const
llvm::DynamicAPInt getIndex() const
void print(mlir::raw_ostream &os) const
IndexRange getIndexRange() const
bool overlaps(const SourceRefIndex &rhs) const
Return true when these path components select any common storage.
component::MemberDefOp getMember() const
SourceRefIndex(component::MemberDefOp f)
llvm::StringRef getPodRecordName() const
SourceRefSet & join(const SourceRefSet &rhs)
A reference to a "source", which is the base value from which other SSA values are derived.
bool isBlockArgument() const
bool overlaps(const SourceRef &rhs) const
Return true when both references select overlapping storage at the same path depth.
mlir::FailureOr< SourceRef > createChild(const SourceRefIndex &r) const
std::vector< SourceRef > getAllChildren(mlir::SymbolTableCollection &tables, mlir::ModuleOp mod) const
Get all direct children of this SourceRef, assuming this ref is not a scalar.
mlir::FailureOr< std::vector< SourceRefIndex > > getSuffix(const SourceRef &prefix) const
If prefix is a valid prefix of this reference, return the suffix that remains after removing the pref...
mlir::FailureOr< function::CallOp > getCallOp() const
void print(mlir::raw_ostream &os) const
Print this reference using source-style names.
bool isCallResult() const
bool operator==(const SourceRef &rhs) const
bool isConstantFelt() const
llvm::ArrayRef< SourceRefIndex > getPath() const
bool isValidPrefix(const SourceRef &prefix) const
Returns true iff prefix is a valid prefix of this reference.
std::strong_ordering operator<=>(const SourceRef &rhs) const
mlir::FailureOr< llvm::DynamicAPInt > getConstantFeltValue() const
bool isConstantIndex() const
std::vector< SourceRefIndex > Path
static std::vector< SourceRef > getAllSourceRefs(mlir::SymbolTableCollection &tables, mlir::ModuleOp mod, const SourceRef &root)
Produce all possible SourceRefs that are present starting from the given root.
mlir::FailureOr< llvm::DynamicAPInt > getConstantValue() const
mlir::FailureOr< unsigned > getInputNum() const
SourceRef narrowRanges(const SourceRef &rhs) const
Return a copy with ranged array indices narrowed by concrete indices from rhs.
mlir::FailureOr< NonDetOp > getNonDetOp() const
mlir::FailureOr< mlir::BlockArgument > getBlockArgument() const
mlir::FailureOr< llvm::DynamicAPInt > getConstantIndexValue() const
mlir::FailureOr< SourceRef > translate(const SourceRef &prefix, const SourceRef &other) const
Create a new reference with prefix replaced with other iff prefix is a valid prefix for this referenc...
bool isTemplateConstant() const
Return whether this reference originates from a template constant read.
bool isConstantInt() const
bool isCreateStructOp() const
mlir::Type getType() const
static constexpr ::llvm::StringLiteral getOperationName()
::llzk::function::FuncDefOp getConstrainFuncOp()
Gets the FuncDefOp that defines the constrain function in this structure, if present,...
::mlir::FailureOr< SymbolLookupResult< StructDefOp > > getDefinition(::mlir::SymbolTableCollection &symbolTable, ::mlir::Operation *op, bool reportMissing=true) const
Gets the struct op that defines this struct.
::mlir::SymbolRefAttr getCallee()
::mlir::Value getSelfValueFromCompute()
Return the "self" value (i.e.
::std::optional<::mlir::StringAttr > getResNameAttr(unsigned index)
Return the function.res_name attribute for the result at the given index.
bool isStructCompute()
Return true iff the function is within a StructDefOp and named FUNC_NAME_COMPUTE.
::llvm::ArrayRef<::llzk::pod::RecordAttr > getRecords() const
ExpressionValue mod(const llvm::SMTSolverRef &solver, const ExpressionValue &lhs, const ExpressionValue &rhs)
void ensure(bool condition, const llvm::Twine &errMsg)
FailureOr< ModuleOp > getRootModule(Operation *from)
ExpressionValue cmp(const llvm::SMTSolverRef &solver, CmpOp op, const ExpressionValue &lhs, const ExpressionValue &rhs)
Interval operator<<(const Interval &lhs, const Interval &rhs)
mlir::FailureOr< SymbolLookupResult< T > > resolveCallable(mlir::SymbolTableCollection &symbolTable, mlir::CallOpInterface call)
Based on mlir::CallOpInterface::resolveCallable, but using LLZK lookup helpers.
std::vector< SourceRef > getAllChildren(SymbolTableCollection &, ModuleOp, ArrayType arrayTy, const SourceRef &root)
mlir::FailureOr< SymbolLookupResultUntyped > lookupSymbolIn(mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, Within &&lookupWithin, mlir::Operation *origin, bool reportMissing=true)
APSInt toAPSInt(const DynamicAPInt &i)
SymbolLookupResult< StructDefOp > getStructDef(SymbolTableCollection &tables, ModuleOp mod, StructType ty)
Lookup a StructDefOp from a given StructType.
Order named operations by source location, using the symbol name to break ties or when source locatio...
size_t operator()(const SourceRefIndex &c) const
size_t operator()(const SourceRef &val) const