1//===- Nodes.cpp ----------------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir/Tools/PDLL/AST/Nodes.h"
10#include "mlir/Tools/PDLL/AST/Context.h"
11#include "llvm/ADT/SmallPtrSet.h"
12#include "llvm/ADT/TypeSwitch.h"
13#include <optional>
14
15using namespace mlir;
16using namespace mlir::pdll::ast;
17
18/// Copy a string reference into the context with a null terminator.
19static StringRef copyStringWithNull(Context &ctx, StringRef str) {
20 if (str.empty())
21 return str;
22
23 char *data = ctx.getAllocator().Allocate<char>(Num: str.size() + 1);
24 std::copy(first: str.begin(), last: str.end(), result: data);
25 data[str.size()] = 0;
26 return StringRef(data, str.size());
27}
28
29//===----------------------------------------------------------------------===//
30// Name
31//===----------------------------------------------------------------------===//
32
33const Name &Name::create(Context &ctx, StringRef name, SMRange location) {
34 return *new (ctx.getAllocator().Allocate<Name>())
35 Name(copyStringWithNull(ctx, str: name), location);
36}
37
38//===----------------------------------------------------------------------===//
39// Node
40//===----------------------------------------------------------------------===//
41
42namespace {
43class NodeVisitor {
44public:
45 explicit NodeVisitor(function_ref<void(const Node *)> visitFn)
46 : visitFn(visitFn) {}
47
48 void visit(const Node *node) {
49 if (!node || !alreadyVisited.insert(Ptr: node).second)
50 return;
51
52 visitFn(node);
53 TypeSwitch<const Node *>(node)
54 .Case<
55 // Statements.
56 const CompoundStmt, const EraseStmt, const LetStmt,
57 const ReplaceStmt, const ReturnStmt, const RewriteStmt,
58
59 // Expressions.
60 const AttributeExpr, const CallExpr, const DeclRefExpr,
61 const MemberAccessExpr, const OperationExpr, const RangeExpr,
62 const TupleExpr, const TypeExpr,
63
64 // Core Constraint Decls.
65 const AttrConstraintDecl, const OpConstraintDecl,
66 const TypeConstraintDecl, const TypeRangeConstraintDecl,
67 const ValueConstraintDecl, const ValueRangeConstraintDecl,
68
69 // Decls.
70 const NamedAttributeDecl, const OpNameDecl, const PatternDecl,
71 const UserConstraintDecl, const UserRewriteDecl, const VariableDecl,
72
73 const Module>(
74 caseFn: [&](auto derivedNode) { this->visitImpl(derivedNode); })
75 .Default(defaultFn: [](const Node *) { llvm_unreachable("unknown AST node"); });
76 }
77
78private:
79 void visitImpl(const CompoundStmt *stmt) {
80 for (const Node *child : stmt->getChildren())
81 visit(node: child);
82 }
83 void visitImpl(const EraseStmt *stmt) { visit(node: stmt->getRootOpExpr()); }
84 void visitImpl(const LetStmt *stmt) { visit(node: stmt->getVarDecl()); }
85 void visitImpl(const ReplaceStmt *stmt) {
86 visit(node: stmt->getRootOpExpr());
87 for (const Node *child : stmt->getReplExprs())
88 visit(node: child);
89 }
90 void visitImpl(const ReturnStmt *stmt) { visit(node: stmt->getResultExpr()); }
91 void visitImpl(const RewriteStmt *stmt) {
92 visit(node: stmt->getRootOpExpr());
93 visit(node: stmt->getRewriteBody());
94 }
95
96 void visitImpl(const AttributeExpr *expr) {}
97 void visitImpl(const CallExpr *expr) {
98 visit(node: expr->getCallableExpr());
99 for (const Node *child : expr->getArguments())
100 visit(node: child);
101 }
102 void visitImpl(const DeclRefExpr *expr) { visit(node: expr->getDecl()); }
103 void visitImpl(const MemberAccessExpr *expr) { visit(node: expr->getParentExpr()); }
104 void visitImpl(const OperationExpr *expr) {
105 visit(node: expr->getNameDecl());
106 for (const Node *child : expr->getOperands())
107 visit(node: child);
108 for (const Node *child : expr->getResultTypes())
109 visit(node: child);
110 for (const Node *child : expr->getAttributes())
111 visit(node: child);
112 }
113 void visitImpl(const RangeExpr *expr) {
114 for (const Node *child : expr->getElements())
115 visit(node: child);
116 }
117 void visitImpl(const TupleExpr *expr) {
118 for (const Node *child : expr->getElements())
119 visit(node: child);
120 }
121 void visitImpl(const TypeExpr *expr) {}
122
123 void visitImpl(const AttrConstraintDecl *decl) { visit(node: decl->getTypeExpr()); }
124 void visitImpl(const OpConstraintDecl *decl) { visit(node: decl->getNameDecl()); }
125 void visitImpl(const TypeConstraintDecl *decl) {}
126 void visitImpl(const TypeRangeConstraintDecl *decl) {}
127 void visitImpl(const ValueConstraintDecl *decl) {
128 visit(node: decl->getTypeExpr());
129 }
130 void visitImpl(const ValueRangeConstraintDecl *decl) {
131 visit(node: decl->getTypeExpr());
132 }
133
134 void visitImpl(const NamedAttributeDecl *decl) { visit(node: decl->getValue()); }
135 void visitImpl(const OpNameDecl *decl) {}
136 void visitImpl(const PatternDecl *decl) { visit(node: decl->getBody()); }
137 void visitImpl(const UserConstraintDecl *decl) {
138 for (const Node *child : decl->getInputs())
139 visit(node: child);
140 for (const Node *child : decl->getResults())
141 visit(node: child);
142 visit(node: decl->getBody());
143 }
144 void visitImpl(const UserRewriteDecl *decl) {
145 for (const Node *child : decl->getInputs())
146 visit(node: child);
147 for (const Node *child : decl->getResults())
148 visit(node: child);
149 visit(node: decl->getBody());
150 }
151 void visitImpl(const VariableDecl *decl) {
152 visit(node: decl->getInitExpr());
153 for (const ConstraintRef &child : decl->getConstraints())
154 visit(node: child.constraint);
155 }
156
157 void visitImpl(const Module *module) {
158 for (const Node *child : module->getChildren())
159 visit(node: child);
160 }
161
162 function_ref<void(const Node *)> visitFn;
163 SmallPtrSet<const Node *, 16> alreadyVisited;
164};
165} // namespace
166
167void Node::walk(function_ref<void(const Node *)> walkFn) const {
168 return NodeVisitor(walkFn).visit(node: this);
169}
170
171//===----------------------------------------------------------------------===//
172// DeclScope
173//===----------------------------------------------------------------------===//
174
175void DeclScope::add(Decl *decl) {
176 const Name *name = decl->getName();
177 assert(name && "expected a named decl");
178 assert(!decls.count(name->getName()) && "decl with this name already exists");
179 decls.try_emplace(Key: name->getName(), Args&: decl);
180}
181
182Decl *DeclScope::lookup(StringRef name) {
183 if (Decl *decl = decls.lookup(Key: name))
184 return decl;
185 return parent ? parent->lookup(name) : nullptr;
186}
187
188//===----------------------------------------------------------------------===//
189// CompoundStmt
190//===----------------------------------------------------------------------===//
191
192CompoundStmt *CompoundStmt::create(Context &ctx, SMRange loc,
193 ArrayRef<Stmt *> children) {
194 unsigned allocSize = CompoundStmt::totalSizeToAlloc<Stmt *>(Counts: children.size());
195 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(CompoundStmt));
196
197 CompoundStmt *stmt = new (rawData) CompoundStmt(loc, children.size());
198 llvm::uninitialized_copy(Src&: children, Dst: stmt->getChildren().begin());
199 return stmt;
200}
201
202//===----------------------------------------------------------------------===//
203// LetStmt
204//===----------------------------------------------------------------------===//
205
206LetStmt *LetStmt::create(Context &ctx, SMRange loc, VariableDecl *varDecl) {
207 return new (ctx.getAllocator().Allocate<LetStmt>()) LetStmt(loc, varDecl);
208}
209
210//===----------------------------------------------------------------------===//
211// OpRewriteStmt
212//===----------------------------------------------------------------------===//
213
214//===----------------------------------------------------------------------===//
215// EraseStmt
216//===----------------------------------------------------------------------===//
217
218EraseStmt *EraseStmt::create(Context &ctx, SMRange loc, Expr *rootOp) {
219 return new (ctx.getAllocator().Allocate<EraseStmt>()) EraseStmt(loc, rootOp);
220}
221
222//===----------------------------------------------------------------------===//
223// ReplaceStmt
224//===----------------------------------------------------------------------===//
225
226ReplaceStmt *ReplaceStmt::create(Context &ctx, SMRange loc, Expr *rootOp,
227 ArrayRef<Expr *> replExprs) {
228 unsigned allocSize = ReplaceStmt::totalSizeToAlloc<Expr *>(Counts: replExprs.size());
229 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(ReplaceStmt));
230
231 ReplaceStmt *stmt = new (rawData) ReplaceStmt(loc, rootOp, replExprs.size());
232 llvm::uninitialized_copy(Src&: replExprs, Dst: stmt->getReplExprs().begin());
233 return stmt;
234}
235
236//===----------------------------------------------------------------------===//
237// RewriteStmt
238//===----------------------------------------------------------------------===//
239
240RewriteStmt *RewriteStmt::create(Context &ctx, SMRange loc, Expr *rootOp,
241 CompoundStmt *rewriteBody) {
242 return new (ctx.getAllocator().Allocate<RewriteStmt>())
243 RewriteStmt(loc, rootOp, rewriteBody);
244}
245
246//===----------------------------------------------------------------------===//
247// ReturnStmt
248//===----------------------------------------------------------------------===//
249
250ReturnStmt *ReturnStmt::create(Context &ctx, SMRange loc, Expr *resultExpr) {
251 return new (ctx.getAllocator().Allocate<ReturnStmt>())
252 ReturnStmt(loc, resultExpr);
253}
254
255//===----------------------------------------------------------------------===//
256// AttributeExpr
257//===----------------------------------------------------------------------===//
258
259AttributeExpr *AttributeExpr::create(Context &ctx, SMRange loc,
260 StringRef value) {
261 return new (ctx.getAllocator().Allocate<AttributeExpr>())
262 AttributeExpr(ctx, loc, copyStringWithNull(ctx, str: value));
263}
264
265//===----------------------------------------------------------------------===//
266// CallExpr
267//===----------------------------------------------------------------------===//
268
269CallExpr *CallExpr::create(Context &ctx, SMRange loc, Expr *callable,
270 ArrayRef<Expr *> arguments, Type resultType,
271 bool isNegated) {
272 unsigned allocSize = CallExpr::totalSizeToAlloc<Expr *>(Counts: arguments.size());
273 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(CallExpr));
274
275 CallExpr *expr = new (rawData)
276 CallExpr(loc, resultType, callable, arguments.size(), isNegated);
277 llvm::uninitialized_copy(Src&: arguments, Dst: expr->getArguments().begin());
278 return expr;
279}
280
281//===----------------------------------------------------------------------===//
282// DeclRefExpr
283//===----------------------------------------------------------------------===//
284
285DeclRefExpr *DeclRefExpr::create(Context &ctx, SMRange loc, Decl *decl,
286 Type type) {
287 return new (ctx.getAllocator().Allocate<DeclRefExpr>())
288 DeclRefExpr(loc, decl, type);
289}
290
291//===----------------------------------------------------------------------===//
292// MemberAccessExpr
293//===----------------------------------------------------------------------===//
294
295MemberAccessExpr *MemberAccessExpr::create(Context &ctx, SMRange loc,
296 const Expr *parentExpr,
297 StringRef memberName, Type type) {
298 return new (ctx.getAllocator().Allocate<MemberAccessExpr>()) MemberAccessExpr(
299 loc, parentExpr, memberName.copy(A&: ctx.getAllocator()), type);
300}
301
302//===----------------------------------------------------------------------===//
303// OperationExpr
304//===----------------------------------------------------------------------===//
305
306OperationExpr *
307OperationExpr::create(Context &ctx, SMRange loc, const ods::Operation *odsOp,
308 const OpNameDecl *name, ArrayRef<Expr *> operands,
309 ArrayRef<Expr *> resultTypes,
310 ArrayRef<NamedAttributeDecl *> attributes) {
311 unsigned allocSize =
312 OperationExpr::totalSizeToAlloc<Expr *, NamedAttributeDecl *>(
313 Counts: operands.size() + resultTypes.size(), Counts: attributes.size());
314 void *rawData =
315 ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(OperationExpr));
316
317 Type resultType = OperationType::get(context&: ctx, name: name->getName(), odsOp);
318 OperationExpr *opExpr = new (rawData)
319 OperationExpr(loc, resultType, name, operands.size(), resultTypes.size(),
320 attributes.size(), name->getLoc());
321 llvm::uninitialized_copy(Src&: operands, Dst: opExpr->getOperands().begin());
322 llvm::uninitialized_copy(Src&: resultTypes, Dst: opExpr->getResultTypes().begin());
323 llvm::uninitialized_copy(Src&: attributes, Dst: opExpr->getAttributes().begin());
324 return opExpr;
325}
326
327std::optional<StringRef> OperationExpr::getName() const {
328 return getNameDecl()->getName();
329}
330
331//===----------------------------------------------------------------------===//
332// RangeExpr
333//===----------------------------------------------------------------------===//
334
335RangeExpr *RangeExpr::create(Context &ctx, SMRange loc,
336 ArrayRef<Expr *> elements, RangeType type) {
337 unsigned allocSize = RangeExpr::totalSizeToAlloc<Expr *>(Counts: elements.size());
338 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(TupleExpr));
339
340 RangeExpr *expr = new (rawData) RangeExpr(loc, type, elements.size());
341 llvm::uninitialized_copy(Src&: elements, Dst: expr->getElements().begin());
342 return expr;
343}
344
345//===----------------------------------------------------------------------===//
346// TupleExpr
347//===----------------------------------------------------------------------===//
348
349TupleExpr *TupleExpr::create(Context &ctx, SMRange loc,
350 ArrayRef<Expr *> elements,
351 ArrayRef<StringRef> names) {
352 unsigned allocSize = TupleExpr::totalSizeToAlloc<Expr *>(Counts: elements.size());
353 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(TupleExpr));
354
355 auto elementTypes = llvm::map_range(
356 C&: elements, F: [](const Expr *expr) { return expr->getType(); });
357 TupleType type = TupleType::get(context&: ctx, elementTypes: llvm::to_vector(Range&: elementTypes), elementNames: names);
358
359 TupleExpr *expr = new (rawData) TupleExpr(loc, type);
360 llvm::uninitialized_copy(Src&: elements, Dst: expr->getElements().begin());
361 return expr;
362}
363
364//===----------------------------------------------------------------------===//
365// TypeExpr
366//===----------------------------------------------------------------------===//
367
368TypeExpr *TypeExpr::create(Context &ctx, SMRange loc, StringRef value) {
369 return new (ctx.getAllocator().Allocate<TypeExpr>())
370 TypeExpr(ctx, loc, copyStringWithNull(ctx, str: value));
371}
372
373//===----------------------------------------------------------------------===//
374// Decl
375//===----------------------------------------------------------------------===//
376
377void Decl::setDocComment(Context &ctx, StringRef comment) {
378 docComment = comment.copy(A&: ctx.getAllocator());
379}
380
381//===----------------------------------------------------------------------===//
382// AttrConstraintDecl
383//===----------------------------------------------------------------------===//
384
385AttrConstraintDecl *AttrConstraintDecl::create(Context &ctx, SMRange loc,
386 Expr *typeExpr) {
387 return new (ctx.getAllocator().Allocate<AttrConstraintDecl>())
388 AttrConstraintDecl(loc, typeExpr);
389}
390
391//===----------------------------------------------------------------------===//
392// OpConstraintDecl
393//===----------------------------------------------------------------------===//
394
395OpConstraintDecl *OpConstraintDecl::create(Context &ctx, SMRange loc,
396 const OpNameDecl *nameDecl) {
397 if (!nameDecl)
398 nameDecl = OpNameDecl::create(ctx, loc: SMRange());
399
400 return new (ctx.getAllocator().Allocate<OpConstraintDecl>())
401 OpConstraintDecl(loc, nameDecl);
402}
403
404std::optional<StringRef> OpConstraintDecl::getName() const {
405 return getNameDecl()->getName();
406}
407
408//===----------------------------------------------------------------------===//
409// TypeConstraintDecl
410//===----------------------------------------------------------------------===//
411
412TypeConstraintDecl *TypeConstraintDecl::create(Context &ctx, SMRange loc) {
413 return new (ctx.getAllocator().Allocate<TypeConstraintDecl>())
414 TypeConstraintDecl(loc);
415}
416
417//===----------------------------------------------------------------------===//
418// TypeRangeConstraintDecl
419//===----------------------------------------------------------------------===//
420
421TypeRangeConstraintDecl *TypeRangeConstraintDecl::create(Context &ctx,
422 SMRange loc) {
423 return new (ctx.getAllocator().Allocate<TypeRangeConstraintDecl>())
424 TypeRangeConstraintDecl(loc);
425}
426
427//===----------------------------------------------------------------------===//
428// ValueConstraintDecl
429//===----------------------------------------------------------------------===//
430
431ValueConstraintDecl *ValueConstraintDecl::create(Context &ctx, SMRange loc,
432 Expr *typeExpr) {
433 return new (ctx.getAllocator().Allocate<ValueConstraintDecl>())
434 ValueConstraintDecl(loc, typeExpr);
435}
436
437//===----------------------------------------------------------------------===//
438// ValueRangeConstraintDecl
439//===----------------------------------------------------------------------===//
440
441ValueRangeConstraintDecl *
442ValueRangeConstraintDecl::create(Context &ctx, SMRange loc, Expr *typeExpr) {
443 return new (ctx.getAllocator().Allocate<ValueRangeConstraintDecl>())
444 ValueRangeConstraintDecl(loc, typeExpr);
445}
446
447//===----------------------------------------------------------------------===//
448// UserConstraintDecl
449//===----------------------------------------------------------------------===//
450
451std::optional<StringRef>
452UserConstraintDecl::getNativeInputType(unsigned index) const {
453 return hasNativeInputTypes ? getTrailingObjects<StringRef>()[index]
454 : std::optional<StringRef>();
455}
456
457UserConstraintDecl *UserConstraintDecl::createImpl(
458 Context &ctx, const Name &name, ArrayRef<VariableDecl *> inputs,
459 ArrayRef<StringRef> nativeInputTypes, ArrayRef<VariableDecl *> results,
460 std::optional<StringRef> codeBlock, const CompoundStmt *body,
461 Type resultType) {
462 bool hasNativeInputTypes = !nativeInputTypes.empty();
463 assert(!hasNativeInputTypes || nativeInputTypes.size() == inputs.size());
464
465 unsigned allocSize =
466 UserConstraintDecl::totalSizeToAlloc<VariableDecl *, StringRef>(
467 Counts: inputs.size() + results.size(),
468 Counts: hasNativeInputTypes ? inputs.size() : 0);
469 void *rawData =
470 ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(UserConstraintDecl));
471 if (codeBlock)
472 codeBlock = codeBlock->copy(A&: ctx.getAllocator());
473
474 UserConstraintDecl *decl = new (rawData)
475 UserConstraintDecl(name, inputs.size(), hasNativeInputTypes,
476 results.size(), codeBlock, body, resultType);
477 llvm::uninitialized_copy(Src&: inputs, Dst: decl->getInputs().begin());
478 llvm::uninitialized_copy(Src&: results, Dst: decl->getResults().begin());
479 if (hasNativeInputTypes) {
480 StringRef *nativeInputTypesPtr = decl->getTrailingObjects<StringRef>();
481 for (unsigned i = 0, e = inputs.size(); i < e; ++i)
482 nativeInputTypesPtr[i] = nativeInputTypes[i].copy(A&: ctx.getAllocator());
483 }
484
485 return decl;
486}
487
488//===----------------------------------------------------------------------===//
489// NamedAttributeDecl
490//===----------------------------------------------------------------------===//
491
492NamedAttributeDecl *NamedAttributeDecl::create(Context &ctx, const Name &name,
493 Expr *value) {
494 return new (ctx.getAllocator().Allocate<NamedAttributeDecl>())
495 NamedAttributeDecl(name, value);
496}
497
498//===----------------------------------------------------------------------===//
499// OpNameDecl
500//===----------------------------------------------------------------------===//
501
502OpNameDecl *OpNameDecl::create(Context &ctx, const Name &name) {
503 return new (ctx.getAllocator().Allocate<OpNameDecl>()) OpNameDecl(name);
504}
505OpNameDecl *OpNameDecl::create(Context &ctx, SMRange loc) {
506 return new (ctx.getAllocator().Allocate<OpNameDecl>()) OpNameDecl(loc);
507}
508
509//===----------------------------------------------------------------------===//
510// PatternDecl
511//===----------------------------------------------------------------------===//
512
513PatternDecl *PatternDecl::create(Context &ctx, SMRange loc, const Name *name,
514 std::optional<uint16_t> benefit,
515 bool hasBoundedRecursion,
516 const CompoundStmt *body) {
517 return new (ctx.getAllocator().Allocate<PatternDecl>())
518 PatternDecl(loc, name, benefit, hasBoundedRecursion, body);
519}
520
521//===----------------------------------------------------------------------===//
522// UserRewriteDecl
523//===----------------------------------------------------------------------===//
524
525UserRewriteDecl *UserRewriteDecl::createImpl(Context &ctx, const Name &name,
526 ArrayRef<VariableDecl *> inputs,
527 ArrayRef<VariableDecl *> results,
528 std::optional<StringRef> codeBlock,
529 const CompoundStmt *body,
530 Type resultType) {
531 unsigned allocSize = UserRewriteDecl::totalSizeToAlloc<VariableDecl *>(
532 Counts: inputs.size() + results.size());
533 void *rawData =
534 ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(UserRewriteDecl));
535 if (codeBlock)
536 codeBlock = codeBlock->copy(A&: ctx.getAllocator());
537
538 UserRewriteDecl *decl = new (rawData) UserRewriteDecl(
539 name, inputs.size(), results.size(), codeBlock, body, resultType);
540 llvm::uninitialized_copy(Src&: inputs, Dst: decl->getInputs().begin());
541 llvm::uninitialized_copy(Src&: results, Dst: decl->getResults().begin());
542 return decl;
543}
544
545//===----------------------------------------------------------------------===//
546// VariableDecl
547//===----------------------------------------------------------------------===//
548
549VariableDecl *VariableDecl::create(Context &ctx, const Name &name, Type type,
550 Expr *initExpr,
551 ArrayRef<ConstraintRef> constraints) {
552 unsigned allocSize =
553 VariableDecl::totalSizeToAlloc<ConstraintRef>(Counts: constraints.size());
554 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(VariableDecl));
555
556 VariableDecl *varDecl =
557 new (rawData) VariableDecl(name, type, initExpr, constraints.size());
558 llvm::uninitialized_copy(Src&: constraints, Dst: varDecl->getConstraints().begin());
559 return varDecl;
560}
561
562//===----------------------------------------------------------------------===//
563// Module
564//===----------------------------------------------------------------------===//
565
566Module *Module::create(Context &ctx, SMLoc loc, ArrayRef<Decl *> children) {
567 unsigned allocSize = Module::totalSizeToAlloc<Decl *>(Counts: children.size());
568 void *rawData = ctx.getAllocator().Allocate(Size: allocSize, Alignment: alignof(Module));
569
570 Module *module = new (rawData) Module(loc, children.size());
571 llvm::uninitialized_copy(Src&: children, Dst: module->getChildren().begin());
572 return module;
573}
574

source code of mlir/lib/Tools/PDLL/AST/Nodes.cpp