LLZK 3.0.0
An open-source IR for Zero Knowledge (ZK) circuits
Loading...
Searching...
No Matches
SourceRef.cpp
Go to the documentation of this file.
1//===-- SourceRef.cpp - SourceRef implementation ----------------*- C++ -*-===//
2//
3// Part of the LLZK Project, under the Apache License v2.0.
4// See LICENSE.txt for license information.
5// Copyright 2025 Veridise Inc.
6// SPDX-License-Identifier: Apache-2.0
7//
8//===----------------------------------------------------------------------===//
9
11
16#include "llzk/Util/Compare.h"
17#include "llzk/Util/Debug.h"
21
22#include <mlir/IR/AsmState.h>
23
24using namespace mlir;
25
26namespace llzk {
27
28using namespace array;
29using namespace component;
30using namespace felt;
31using namespace function;
32using namespace polymorphic;
33using namespace string;
34
35namespace {
36
37std::strong_ordering
38compareDynamicAPInt(const llvm::DynamicAPInt &lhs, const llvm::DynamicAPInt &rhs) {
39 if (lhs < rhs) {
40 return std::strong_ordering::less;
41 }
42 if (rhs < lhs) {
43 return std::strong_ordering::greater;
44 }
45 return std::strong_ordering::equal;
46}
47
48std::strong_ordering compareStringRef(llvm::StringRef lhs, llvm::StringRef rhs) {
49 int cmp = lhs.compare(rhs);
50 if (cmp < 0) {
51 return std::strong_ordering::less;
52 }
53 if (cmp > 0) {
54 return std::strong_ordering::greater;
55 }
56 return std::strong_ordering::equal;
57}
58
59// Prints SourceRef path using source-style syntax, i.e. `.` for struct members
60// and pod records and `[...]` for array indices.
61// Returns failure if an unexpected `SourceRefIndex` type is encountered.
62LogicalResult printSourceStylePath(raw_ostream &os, llvm::ArrayRef<SourceRefIndex> path) {
63 for (const auto &idx : path) {
64 if (idx.isMember()) {
65 os << '.' << idx.getMember().getName();
66 continue;
67 }
68 if (idx.isPodRecord()) {
69 os << '.' << idx.getPodRecordName();
70 continue;
71 }
72 if (idx.isIndex() || idx.isIndexRange()) {
73 os << '[' << idx << ']';
74 continue;
75 }
76 return failure();
77 }
78 return success();
79}
80
81std::strong_ordering
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) {
85 return cmp;
86 }
87 }
88 return lhs.size() <=> rhs.size();
89}
90
91} // namespace
92
93/* SourceRefIndex */
94
95void SourceRefIndex::print(raw_ostream &os) const {
96 if (isMember()) {
97 os << '@' << getMember().getName();
98 } else if (isPodRecord()) {
99 os << '@' << getPodRecordName();
100 } else if (isIndex()) {
101 os << getIndex();
102 } else {
103 auto [low, high] = getIndexRange();
104 if (ShapedType::isDynamic(int64_t(high))) {
105 os << "<dynamic>";
106 } else {
107 os << low << ':' << high;
108 }
109 }
110}
111
112std::strong_ordering SourceRefIndex::operator<=>(const SourceRefIndex &rhs) const {
113 if (isMember() && rhs.isMember()) {
115 return std::strong_ordering::less;
116 }
118 return std::strong_ordering::greater;
119 }
120 return std::strong_ordering::equal;
121 }
122 if (isPodRecord() && rhs.isPodRecord()) {
123 return compareStringRef(getPodRecordName(), rhs.getPodRecordName());
124 }
125 if (isIndex() && rhs.isIndex()) {
126 return compareDynamicAPInt(getIndex(), rhs.getIndex());
127 }
128 if (isIndexRange() && rhs.isIndexRange()) {
129 auto [ll, lu] = getIndexRange();
130 auto [rl, ru] = rhs.getIndexRange();
131 if (auto cmp = compareDynamicAPInt(ll, rl); cmp != std::strong_ordering::equal) {
132 return cmp;
133 }
134 return compareDynamicAPInt(lu, ru);
135 }
136
137 if (isMember()) {
138 return std::strong_ordering::less;
139 }
140 if (rhs.isMember()) {
141 return std::strong_ordering::greater;
142 }
143 if (isPodRecord()) {
144 return std::strong_ordering::less;
145 }
146 if (rhs.isPodRecord()) {
147 return std::strong_ordering::greater;
148 }
149 if (isIndex()) {
150 return std::strong_ordering::less;
151 }
152 return std::strong_ordering::greater;
153}
154
156 if (c.isIndex()) {
157 // We don't hash the index directly, because the built-in LLVM hash includes
158 // the bitwidth of the APInt in the hash, which is undesirable for this application.
159 // i.e., We want a N-bit version of x to hash to the same value as an M-bit version of X,
160 // because our equality checks would consider them equal regardless of bitwidth.
161 APSInt idx = toAPSInt(c.getIndex());
162 unsigned requiredBits = idx.getSignificantBits();
163 auto hash = llvm::hash_value(idx.trunc(requiredBits));
164 return hash;
165 } else if (c.isIndexRange()) {
166 auto r = c.getIndexRange();
167 return llvm::hash_value(std::get<0>(r)) ^ llvm::hash_value(std::get<1>(r));
168 } else if (c.isPodRecord()) {
169 return llvm::hash_value(c.getPodRecordName());
170 } else {
172 }
173}
174
176 if (isIndex() && rhs.isIndexRange()) {
177 auto [low, high] = rhs.getIndexRange();
178 return low <= getIndex() && getIndex() < high;
179 }
180 if (isIndexRange() && rhs.isIndex()) {
181 return rhs.overlaps(*this);
182 }
183 if (isIndexRange() && rhs.isIndexRange()) {
184 auto [lhsLow, lhsHigh] = getIndexRange();
185 auto [rhsLow, rhsHigh] = rhs.getIndexRange();
186 return lhsLow < rhsHigh && rhsLow < lhsHigh;
187 }
188 return *this == rhs;
189}
190
191/* SourceRef */
192
193SourceRef::SortCategory SourceRef::getSortCategory() const {
194 if (isBlockArgument()) {
195 return SortCategory::BlockArgument;
196 }
197 if (isCreateStructOp()) {
198 return SortCategory::CreateStruct;
199 }
200 if (isNonDetOp()) {
201 return SortCategory::NonDet;
202 }
203 if (isRooted()) {
204 return SortCategory::RootResult;
205 }
206 if (isTemplateConstant()) {
207 return SortCategory::TemplateConstant;
208 }
209 if (isConstantIndex()) {
210 return SortCategory::ConstantIndex;
211 }
212 if (isConstantFelt()) {
213 return SortCategory::ConstantFelt;
214 }
215
216 llvm::errs() << *this << '\n';
217 llvm_unreachable("unhandled SourceRef sort category");
218}
219
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();
226}
227
228std::strong_ordering
229SourceRef::compareWithinCategory(const SourceRef &rhs, SortCategory category) const {
230 switch (category) {
231 case SortCategory::BlockArgument: {
232 if (auto cmp = *getInputNum() <=> *rhs.getInputNum(); cmp != std::strong_ordering::equal) {
233 return cmp;
234 }
235 if (auto cmp = getAsOpaquePointer() <=> rhs.getAsOpaquePointer();
236 cmp != std::strong_ordering::equal) {
237 return cmp;
238 }
239 return compareSourceRefPaths(getPath(), rhs.getPath());
240 }
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) {
246 return cmp;
247 }
248 return compareSourceRefPaths(getPath(), rhs.getPath());
249 }
250 case SortCategory::TemplateConstant: {
251 if (auto cmp = compareStringRef(getTemplateConstantName(), rhs.getTemplateConstantName());
252 cmp != std::strong_ordering::equal) {
253 return cmp;
254 }
255 return getAsOpaquePointer() <=> rhs.getAsOpaquePointer();
256 }
257 case SortCategory::ConstantIndex:
258 return compareDynamicAPInt(*getConstantIndexValue(), *rhs.getConstantIndexValue());
259 case SortCategory::ConstantFelt:
260 return compareDynamicAPInt(*getConstantFeltValue(), *rhs.getConstantFeltValue());
261 }
262
263 llvm_unreachable("unhandled SourceRef category compare");
264}
265
273SymbolLookupResult<StructDefOp>
274getStructDef(SymbolTableCollection &tables, ModuleOp mod, StructType ty) {
275 auto sDef = ty.getDefinition(tables, mod);
276 ensure(
277 succeeded(sDef),
278 "could not find '" + StructDefOp::getOperationName() + "' op from struct type"
279 );
280
281 return std::move(*sDef);
282}
283
284std::vector<SourceRef>
285SourceRef::getAllSourceRefs(SymbolTableCollection &tables, ModuleOp mod, const SourceRef &root) {
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());
290 }
291 return res;
292}
293
294std::vector<SourceRef> SourceRef::getAllSourceRefs(StructDefOp structDef, FuncDefOp fnOp) {
295 std::vector<SourceRef> res;
296
297 ensure(
298 structDef == fnOp->getParentOfType<StructDefOp>(), "function must be within the given struct"
299 );
300
301 FailureOr<ModuleOp> modOp = getRootModule(structDef);
302 ensure(succeeded(modOp), "could not lookup module from struct " + Twine(structDef.getName()));
303
304 SymbolTableCollection tables;
305 for (auto a : fnOp.getArguments()) {
306 auto argRes = getAllSourceRefs(tables, modOp.value(), SourceRef(a));
307 res.insert(res.end(), argRes.begin(), argRes.end());
308 }
309
310 // For compute functions, the "self" member is not arg0 like for constrain, but
311 // rather the struct value returned from the function.
312 if (fnOp.isStructCompute()) {
313 Value selfVal = fnOp.getSelfValueFromCompute();
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());
318 }
319
320 return res;
321}
322
323std::vector<SourceRef> SourceRef::getAllSourceRefs(StructDefOp structDef, MemberDefOp memberDef) {
324 std::vector<SourceRef> res;
325 FuncDefOp constrainFnOp = structDef.getConstrainFuncOp();
326 ensure(
327 memberDef->getParentOfType<StructDefOp>() == structDef,
328 "Member " + Twine(memberDef.getName()) + " is not a member of struct " +
329 Twine(structDef.getName())
330 );
331 FailureOr<ModuleOp> modOp = getRootModule(structDef);
332 ensure(succeeded(modOp), "could not lookup module from struct " + Twine(structDef.getName()));
333
334 // Get the self argument (like `FuncDefOp::getSelfValueFromConstrain()`)
335 BlockArgument self = constrainFnOp.getArguments().front();
336 SourceRef memberRef = SourceRef(self, {SourceRefIndex(memberDef)});
337
338 SymbolTableCollection tables;
339 return getAllSourceRefs(tables, modOp.value(), memberRef);
340}
341
342Type SourceRef::getType() const {
343 Type currTy = value.getType();
344 for (const auto &idx : getPath()) {
345 if (idx.isMember()) {
346 currTy = idx.getMember().getType();
347 continue;
348 }
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"
355 );
356 });
357 ensure(succeeded(lookup), "SourceRef references a missing pod record");
358 currTy = *lookup;
359 continue;
360 }
361
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);
365 }
366 return currTy;
367}
368
369bool SourceRef::isValidPrefix(const SourceRef &prefix) const {
370 if (isConstant() || prefix.isConstant()) {
371 return false;
372 }
373
374 auto pathRef = getPath();
375 auto prefixPath = prefix.getPath();
376 if (value != prefix.value || pathRef.size() < prefixPath.size()) {
377 return false;
378 }
379 for (size_t i = 0; i < prefixPath.size(); i++) {
380 if (pathRef[i] != prefixPath[i]) {
381 return false;
382 }
383 }
384 return true;
385}
386
387bool SourceRef::overlaps(const SourceRef &rhs) const {
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) {
392 return nullptr;
393 }
394 return func->getParentOfType<StructDefOp>();
395 }
396 auto blockArg = ref.getBlockArgument();
397 if (failed(blockArg)) {
398 return nullptr;
399 }
400 auto func = dyn_cast_if_present<FuncDefOp>(blockArg->getOwner()->getParentOp());
401 return func && func.isStructConstrain() && func.getSelfValueFromConstrain() == *blockArg
402 ? func->getParentOfType<StructDefOp>()
403 : nullptr;
404 };
405 bool sameRoot = value == rhs.value;
406 if (!sameRoot) {
407 StructDefOp lhsStruct = getSelfStruct(*this);
408 StructDefOp rhsStruct = getSelfStruct(rhs);
409 sameRoot = lhsStruct && lhsStruct == rhsStruct;
410 }
411 if (isConstant() || rhs.isConstant() || !sameRoot || path.size() != rhs.path.size()) {
412 return false;
413 }
414 return llvm::all_of(llvm::zip(path, rhs.path), [](const auto &indices) {
415 return std::get<0>(indices).overlaps(std::get<1>(indices));
416 });
417}
418
419SourceRef SourceRef::narrowRanges(const SourceRef &rhs) const {
420 llvm::SmallVector<SourceRefIndex> selections;
421 llvm::copy_if(rhs.getPath(), std::back_inserter(selections), [](const SourceRefIndex &index) {
422 return index.isIndex() || index.isIndexRange();
423 });
424
425 SourceRef result = *this;
426 size_t dimension = 0;
427 for (SourceRefIndex &index : result.getPathMut()) {
428 if (!index.isIndex() && !index.isIndexRange()) {
429 continue;
430 }
431 if (dimension < selections.size() && index.isIndexRange() && selections[dimension].isIndex() &&
432 index.overlaps(selections[dimension])) {
433 index = selections[dimension];
434 }
435 ++dimension;
436 }
437 return result;
438}
439
440FailureOr<SourceRef::Path> SourceRef::getSuffix(const SourceRef &prefix) const {
441 if (!isValidPrefix(prefix)) {
442 return failure();
443 }
444 Path suffix;
445 auto pathRef = getPath();
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]);
450 }
451 return suffix;
452}
453
454FailureOr<SourceRef> SourceRef::translate(const SourceRef &prefix, const SourceRef &other) const {
455 if (isConstant()) {
456 return *this;
457 }
458 auto suffix = getSuffix(prefix);
459 if (failed(suffix)) {
460 return failure();
461 }
462
463 SourceRef newSignalUsage = other; // copy
464 if (newSignalUsage.isRooted()) {
465 SourceRef::Path &pathRef = newSignalUsage.getPathMut();
466 pathRef.insert(pathRef.end(), suffix->begin(), suffix->end());
467 }
468
469 return newSignalUsage;
470}
471
472std::vector<SourceRef> getAllChildren(
473 SymbolTableCollection & /*tables*/, ModuleOp /*mod*/, ArrayType arrayTy, const SourceRef &root
474) {
475 std::vector<SourceRef> res;
476 // Recurse into arrays by iterating over their elements
477 for (int64_t i = 0; i < arrayTy.getDimSize(0); i++) {
478 auto childRef = root.createChild(SourceRefIndex(i));
479 ensure(succeeded(childRef), "array children require a rooted SourceRef");
480 res.push_back(*childRef);
481 }
482
483 return res;
484}
485
486std::vector<SourceRef> getAllChildren(
487 SymbolTableCollection &tables, ModuleOp mod, SymbolLookupResult<StructDefOp> structDefRes,
488 const SourceRef &root
489) {
490 std::vector<SourceRef> res;
491 // Recurse into struct types by iterating over all their member definitions
492 for (auto f : structDefRes.get().getOps<MemberDefOp>()) {
493 // We want to store the MemberDefOp, but without the possibility of accidentally dropping the
494 // reference, so we need to re-lookup the symbol to create a SymbolLookupResult, which will
495 // manage the external module containing the member defs, if needed.
496 // TODO: It would be nice if we could manage module op references differently
497 // so we don't have to do this.
498 auto structDefCopy = structDefRes;
499 auto memberLookup = lookupSymbolIn<MemberDefOp>(
500 tables, SymbolRefAttr::get(f.getContext(), f.getSymNameAttr()), std::move(structDefCopy),
501 mod.getOperation()
502 );
503 ensure(succeeded(memberLookup), "could not get SymbolLookupResult of existing MemberDefOp");
504 auto childRef = root.createChild(SourceRefIndex(memberLookup.value()));
505 ensure(succeeded(childRef), "struct children require a rooted SourceRef");
506 // Make a reference to the current member, regardless of if it is a composite
507 // type or not.
508 res.push_back(*childRef);
509 }
510 return res;
511}
512
513std::vector<SourceRef> getAllChildren(pod::PodType podTy, const SourceRef &root) {
514 std::vector<SourceRef> res;
515 for (auto record : podTy.getRecords()) {
516 auto childRef = root.createChild(SourceRefIndex(record.getName()));
517 ensure(succeeded(childRef), "pod children require a rooted SourceRef");
518 res.push_back(*childRef);
519 }
520 return res;
521}
522
523std::vector<SourceRef>
524SourceRef::getAllChildren(SymbolTableCollection &tables, ModuleOp mod) const {
525 auto ty = getType();
526 if (auto structTy = dyn_cast<StructType>(ty)) {
527 return llzk::getAllChildren(tables, mod, getStructDef(tables, mod, structTy), *this);
528 } else if (auto podTy = dyn_cast<pod::PodType>(ty)) {
529 return llzk::getAllChildren(podTy, *this);
530 } else if (auto arrayType = dyn_cast<ArrayType>(ty)) {
531 return llzk::getAllChildren(tables, mod, arrayType, *this);
532 }
533 // Scalar type, no children
534 return {};
535}
536
537static void printCallResultFallback(raw_ostream &os, function::CallOp callOp, Value value) {
538 os << "<call " << callOp.getCallee();
539 os << ' ';
540 Operation *printScope = callOp.getOperation();
541 if (auto funcOp = callOp->getParentOfType<FuncDefOp>()) {
542 printScope = funcOp.getOperation();
543 }
544 // Allows us to print the SSA result value of the call to disambiguate
545 // repeated calls in the same function.
546 AsmState state(printScope);
547 value.printAsOperand(os, state);
548 os << '>';
549}
550
551static bool shouldPrintNamedCallResult(
552 function::CallOp callOp, OpResult callResult, function::FuncDefOp calleeFunc
553) {
554 auto resName = calleeFunc.getResNameAttr(callResult.getResultNumber());
555 if (!resName) {
556 return false;
557 }
558
559 auto parentFunc = callOp->getParentOfType<FuncDefOp>();
560 if (!parentFunc) {
561 return true;
562 }
563
564 bool foundThisCall = false;
565 bool foundDuplicate = false;
566 parentFunc.walk([&](function::CallOp otherCall) {
567 if (foundDuplicate) {
568 return WalkResult::interrupt();
569 }
570
571 auto otherFunc = llvm::dyn_cast_if_present<FuncDefOp>(otherCall.resolveCallable());
572 if (!otherFunc) {
573 return WalkResult::advance();
574 }
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()) {
579 continue;
580 }
581 if (otherResult == callResult) {
582 foundThisCall = true;
583 continue;
584 }
585 foundDuplicate = true;
586 return WalkResult::interrupt();
587 }
588 return WalkResult::advance();
589 });
590
591 return foundThisCall && !foundDuplicate;
592}
593
594void SourceRef::print(raw_ostream &os) const {
595 if (isConstantFelt()) {
596 os << "<felt.const: " << *getConstantFeltValue() << '>';
597 } else if (isConstantIndex()) {
598 os << "<index: " << *getConstantIndexValue() << '>';
599 } else if (isTemplateConstant()) {
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() << "]>";
605 } else {
606 if (isCreateStructOp()) {
607 os << "%self";
608 } else if (isBlockArgument()) {
609 auto blockArg = *getBlockArgument();
610 auto funcOp = llvm::dyn_cast<FuncDefOp>(blockArg.getOwner()->getParentOp());
611 // The entry self argument of a struct constrain function matches the struct value
612 // returned by the corresponding compute function.
613 if (funcOp && funcOp.isStructConstrain() && funcOp.getSelfValueFromConstrain() == blockArg) {
614 os << "%self";
615 } else {
616 std::optional<StringAttr> argName;
617 if (funcOp) {
618 argName = funcOp.getArgNameAttr(blockArg.getArgNumber());
619 }
620 if (argName) {
621 os << argName->getValue();
622 } else {
623 os << "%arg" << *getInputNum();
624 }
625 }
626 } else if (isNonDetOp()) {
627 os << '<' << *getNonDetOp() << '>';
628 } else if (isCallResult()) {
629 auto callOp = *getCallOp();
630 auto callResult = llvm::cast<OpResult>(value);
631 auto callee = resolveCallable<FuncDefOp>(callOp);
632 if (succeeded(callee)) {
633 FuncDefOp calleeFunc = (*callee).get();
634 if (calleeFunc && shouldPrintNamedCallResult(callOp, callResult, calleeFunc)) {
635 auto resName = *calleeFunc.getResNameAttr(callResult.getResultNumber());
636 os << resName.getValue();
637 } else {
638 printCallResultFallback(os, callOp, value);
639 }
640 } else {
641 printCallResultFallback(os, callOp, value);
642 }
643 } else {
644 ensure(isRooted(), "unhandled print case");
645 OpPrintingFlags flags;
646 value.printAsOperand(os, flags);
647 }
648
649 auto res = printSourceStylePath(os, getPath());
650 ensure(succeeded(res), "unhandled path print case");
651 }
652}
653
654bool SourceRef::operator==(const SourceRef &rhs) const {
655 // This way two felt constants can be equal even if the declared in separate ops.
656 if (isConstantInt() && rhs.isConstantInt()) {
657 DynamicAPInt lhsVal = *getConstantValue(), rhsVal = *rhs.getConstantValue();
658 return getType() == rhs.getType() && lhsVal == rhsVal;
659 }
660 return constant == rhs.constant && value == rhs.value && llvm::equal(getPath(), rhs.getPath());
661}
662
663// required for EquivalenceClasses usage
664std::strong_ordering SourceRef::operator<=>(const SourceRef &rhs) const {
665 auto lhsCategory = getSortCategory();
666 auto rhsCategory = rhs.getSortCategory();
667 if (auto cmp = lhsCategory <=> rhsCategory; cmp != std::strong_ordering::equal) {
668 return cmp;
669 }
670 return compareWithinCategory(rhs, lhsCategory);
671}
672
673size_t SourceRef::Hash::operator()(const SourceRef &val) const {
674 if (val.isConstantInt()) {
675 return llvm::hash_combine(val.getType(), *val.getConstantValue());
676 } else if (val.isTemplateConstant()) {
677 return llvm::hash_value(val.getAsOpaquePointer());
678 } else {
679 ensure(
680 val.isBlockArgument() || val.isCreateStructOp() || val.isNonDetOp() || val.isRooted(),
681 "unhandled SourceRef hash case"
682 );
683
684 size_t hash = llvm::hash_value(val.getAsOpaquePointer());
685 for (const auto &f : val.getPath()) {
686 hash = llvm::hash_combine(hash, f.getHash());
687 }
688 return hash;
689 }
690}
691
692raw_ostream &operator<<(raw_ostream &os, const SourceRef &rhs) {
693 rhs.print(os);
694 return os;
695}
696
697/* SourceRefSet */
698
700 insert(rhs.begin(), rhs.end());
701 return *this;
702}
703
704raw_ostream &operator<<(raw_ostream &os, const SourceRefSet &rhs) {
705 os << "{ ";
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();) {
709 os << *it;
710 it++;
711 if (it != sortedRefs.end()) {
712 os << ", ";
713 } else {
714 os << ' ';
715 }
716 }
717 os << '}';
718 return os;
719}
720
721} // namespace llzk
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.
Definition SourceRef.h:43
std::strong_ordering operator<=>(const SourceRefIndex &rhs) const
bool isIndexRange() const
Definition SourceRef.h:82
bool isIndex() const
Definition SourceRef.h:76
bool isMember() const
Definition SourceRef.h:57
bool isPodRecord() const
Definition SourceRef.h:69
llvm::DynamicAPInt getIndex() const
Definition SourceRef.h:77
void print(mlir::raw_ostream &os) const
Definition SourceRef.cpp:95
IndexRange getIndexRange() const
Definition SourceRef.h:83
bool overlaps(const SourceRefIndex &rhs) const
Return true when these path components select any common storage.
component::MemberDefOp getMember() const
Definition SourceRef.h:61
SourceRefIndex(component::MemberDefOp f)
Definition SourceRef.h:47
llvm::StringRef getPodRecordName() const
Definition SourceRef.h:74
SourceRefSet & join(const SourceRefSet &rhs)
A reference to a "source", which is the base value from which other SSA values are derived.
Definition SourceRef.h:146
bool isBlockArgument() const
Definition SourceRef.h:260
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
Definition SourceRef.h:368
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
Definition SourceRef.h:296
void print(mlir::raw_ostream &os) const
Print this reference using source-style names.
bool isCallResult() const
Definition SourceRef.h:295
bool operator==(const SourceRef &rhs) const
bool isConstantFelt() const
Definition SourceRef.h:235
bool isRooted() const
Definition SourceRef.h:258
llvm::ArrayRef< SourceRefIndex > getPath() const
Definition SourceRef.h:390
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
Definition SourceRef.h:298
bool isConstantIndex() const
Definition SourceRef.h:238
std::vector< SourceRefIndex > Path
Definition SourceRef.h:148
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
Definition SourceRef.h:313
mlir::FailureOr< unsigned > getInputNum() const
Definition SourceRef.h:279
SourceRef narrowRanges(const SourceRef &rhs) const
Return a copy with ranged array indices narrowed by concrete indices from rhs.
mlir::FailureOr< NonDetOp > getNonDetOp() const
Definition SourceRef.h:293
mlir::FailureOr< mlir::BlockArgument > getBlockArgument() const
Definition SourceRef.h:273
mlir::FailureOr< llvm::DynamicAPInt > getConstantIndexValue() const
Definition SourceRef.h:306
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 isNonDetOp() const
Definition SourceRef.h:292
bool isTemplateConstant() const
Return whether this reference originates from a template constant read.
Definition SourceRef.h:243
bool isConstant() const
Definition SourceRef.h:247
bool isConstantInt() const
Definition SourceRef.h:248
bool isCreateStructOp() const
Definition SourceRef.h:287
mlir::Type getType() const
static constexpr ::llvm::StringLiteral getOperationName()
Definition Ops.h.inc:1170
::llzk::function::FuncDefOp getConstrainFuncOp()
Gets the FuncDefOp that defines the constrain function in this structure, if present,...
Definition Ops.cpp:468
::mlir::FailureOr< SymbolLookupResult< StructDefOp > > getDefinition(::mlir::SymbolTableCollection &symbolTable, ::mlir::Operation *op, bool reportMissing=true) const
Gets the struct op that defines this struct.
Definition Types.cpp:26
::mlir::SymbolRefAttr getCallee()
Definition Ops.cpp.inc:470
::mlir::Value getSelfValueFromCompute()
Return the "self" value (i.e.
Definition Ops.cpp:481
::std::optional<::mlir::StringAttr > getResNameAttr(unsigned index)
Return the function.res_name attribute for the result at the given index.
Definition Ops.cpp:333
bool isStructCompute()
Return true iff the function is within a StructDefOp and named FUNC_NAME_COMPUTE.
Definition Ops.h.inc:912
::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...
Definition Compare.h:75
size_t operator()(const SourceRefIndex &c) const
size_t operator()(const SourceRef &val) const