1 //===- AffineParser.cpp - MLIR Affine Parser ------------------------------===//
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 // This file implements a parser for Affine structures.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "Parser.h"
14 #include "mlir/IR/AffineMap.h"
15 #include "mlir/IR/IntegerSet.h"
16 #include "llvm/Support/SourceMgr.h"
17 
18 using namespace mlir;
19 using namespace mlir::detail;
20 
21 namespace {
22 
23 /// Lower precedence ops (all at the same precedence level). LNoOp is false in
24 /// the boolean sense.
25 enum AffineLowPrecOp {
26   /// Null value.
27   LNoOp,
28   Add,
29   Sub
30 };
31 
32 /// Higher precedence ops - all at the same precedence level. HNoOp is false
33 /// in the boolean sense.
34 enum AffineHighPrecOp {
35   /// Null value.
36   HNoOp,
37   Mul,
38   FloorDiv,
39   CeilDiv,
40   Mod
41 };
42 
43 /// This is a specialized parser for affine structures (affine maps, affine
44 /// expressions, and integer sets), maintaining the state transient to their
45 /// bodies.
46 class AffineParser : public Parser {
47 public:
48   AffineParser(ParserState &state, bool allowParsingSSAIds = false,
49                function_ref<ParseResult(bool)> parseElement = nullptr)
50       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
51         parseElement(parseElement) {}
52 
53   ParseResult parseAffineMapRange(unsigned numDims, unsigned numSymbols,
54                                   AffineMap &result);
55   ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set);
56   ParseResult parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols,
57                                          IntegerSet &result);
58   ParseResult parseAffineMapOfSSAIds(AffineMap &map,
59                                      OpAsmParser::Delimiter delimiter);
60   ParseResult parseAffineExprOfSSAIds(AffineExpr &expr);
61   void getDimsAndSymbolSSAIds(SmallVectorImpl<StringRef> &dimAndSymbolSSAIds,
62                               unsigned &numDims);
63 
64 private:
65   // Binary affine op parsing.
66   AffineLowPrecOp consumeIfLowPrecOp();
67   AffineHighPrecOp consumeIfHighPrecOp();
68 
69   // Identifier lists for polyhedral structures.
70   ParseResult parseDimIdList(unsigned &numDims);
71   ParseResult parseSymbolIdList(unsigned &numSymbols);
72   ParseResult parseDimAndOptionalSymbolIdList(unsigned &numDims,
73                                               unsigned &numSymbols);
74   ParseResult parseIdentifierDefinition(AffineExpr idExpr);
75 
76   AffineExpr parseAffineExpr();
77   AffineExpr parseParentheticalExpr();
78   AffineExpr parseNegateExpression(AffineExpr lhs);
79   AffineExpr parseIntegerExpr();
80   AffineExpr parseBareIdExpr();
81   AffineExpr parseSSAIdExpr(bool isSymbol);
82   AffineExpr parseSymbolSSAIdExpr();
83 
84   AffineExpr getAffineBinaryOpExpr(AffineHighPrecOp op, AffineExpr lhs,
85                                    AffineExpr rhs, SMLoc opLoc);
86   AffineExpr getAffineBinaryOpExpr(AffineLowPrecOp op, AffineExpr lhs,
87                                    AffineExpr rhs);
88   AffineExpr parseAffineOperandExpr(AffineExpr lhs);
89   AffineExpr parseAffineLowPrecOpExpr(AffineExpr llhs, AffineLowPrecOp llhsOp);
90   AffineExpr parseAffineHighPrecOpExpr(AffineExpr llhs, AffineHighPrecOp llhsOp,
91                                        SMLoc llhsOpLoc);
92   AffineExpr parseAffineConstraint(bool *isEq);
93 
94 private:
95   bool allowParsingSSAIds;
96   function_ref<ParseResult(bool)> parseElement;
97   unsigned numDimOperands = 0;
98   unsigned numSymbolOperands = 0;
99   SmallVector<std::pair<StringRef, AffineExpr>, 4> dimsAndSymbols;
100 };
101 } // namespace
102 
103 /// Create an affine binary high precedence op expression (mul's, div's, mod).
104 /// opLoc is the location of the op token to be used to report errors
105 /// for non-conforming expressions.
106 AffineExpr AffineParser::getAffineBinaryOpExpr(AffineHighPrecOp op,
107                                                AffineExpr lhs, AffineExpr rhs,
108                                                SMLoc opLoc) {
109   // TODO: make the error location info accurate.
110   switch (op) {
111   case Mul:
112     if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant()) {
113       emitError(opLoc, "non-affine expression: at least one of the multiply "
114                        "operands has to be either a constant or symbolic");
115       return nullptr;
116     }
117     return lhs * rhs;
118   case FloorDiv:
119     if (!rhs.isSymbolicOrConstant()) {
120       emitError(opLoc, "non-affine expression: right operand of floordiv "
121                        "has to be either a constant or symbolic");
122       return nullptr;
123     }
124     return lhs.floorDiv(rhs);
125   case CeilDiv:
126     if (!rhs.isSymbolicOrConstant()) {
127       emitError(opLoc, "non-affine expression: right operand of ceildiv "
128                        "has to be either a constant or symbolic");
129       return nullptr;
130     }
131     return lhs.ceilDiv(rhs);
132   case Mod:
133     if (!rhs.isSymbolicOrConstant()) {
134       emitError(opLoc, "non-affine expression: right operand of mod "
135                        "has to be either a constant or symbolic");
136       return nullptr;
137     }
138     return lhs % rhs;
139   case HNoOp:
140     llvm_unreachable("can't create affine expression for null high prec op");
141     return nullptr;
142   }
143   llvm_unreachable("Unknown AffineHighPrecOp");
144 }
145 
146 /// Create an affine binary low precedence op expression (add, sub).
147 AffineExpr AffineParser::getAffineBinaryOpExpr(AffineLowPrecOp op,
148                                                AffineExpr lhs, AffineExpr rhs) {
149   switch (op) {
150   case AffineLowPrecOp::Add:
151     return lhs + rhs;
152   case AffineLowPrecOp::Sub:
153     return lhs - rhs;
154   case AffineLowPrecOp::LNoOp:
155     llvm_unreachable("can't create affine expression for null low prec op");
156     return nullptr;
157   }
158   llvm_unreachable("Unknown AffineLowPrecOp");
159 }
160 
161 /// Consume this token if it is a lower precedence affine op (there are only
162 /// two precedence levels).
163 AffineLowPrecOp AffineParser::consumeIfLowPrecOp() {
164   switch (getToken().getKind()) {
165   case Token::plus:
166     consumeToken(Token::plus);
167     return AffineLowPrecOp::Add;
168   case Token::minus:
169     consumeToken(Token::minus);
170     return AffineLowPrecOp::Sub;
171   default:
172     return AffineLowPrecOp::LNoOp;
173   }
174 }
175 
176 /// Consume this token if it is a higher precedence affine op (there are only
177 /// two precedence levels)
178 AffineHighPrecOp AffineParser::consumeIfHighPrecOp() {
179   switch (getToken().getKind()) {
180   case Token::star:
181     consumeToken(Token::star);
182     return Mul;
183   case Token::kw_floordiv:
184     consumeToken(Token::kw_floordiv);
185     return FloorDiv;
186   case Token::kw_ceildiv:
187     consumeToken(Token::kw_ceildiv);
188     return CeilDiv;
189   case Token::kw_mod:
190     consumeToken(Token::kw_mod);
191     return Mod;
192   default:
193     return HNoOp;
194   }
195 }
196 
197 /// Parse a high precedence op expression list: mul, div, and mod are high
198 /// precedence binary ops, i.e., parse a
199 ///   expr_1 op_1 expr_2 op_2 ... expr_n
200 /// where op_1, op_2 are all a AffineHighPrecOp (mul, div, mod).
201 /// All affine binary ops are left associative.
202 /// Given llhs, returns (llhs llhsOp lhs) op rhs, or (lhs op rhs) if llhs is
203 /// null. If no rhs can be found, returns (llhs llhsOp lhs) or lhs if llhs is
204 /// null. llhsOpLoc is the location of the llhsOp token that will be used to
205 /// report an error for non-conforming expressions.
206 AffineExpr AffineParser::parseAffineHighPrecOpExpr(AffineExpr llhs,
207                                                    AffineHighPrecOp llhsOp,
208                                                    SMLoc llhsOpLoc) {
209   AffineExpr lhs = parseAffineOperandExpr(llhs);
210   if (!lhs)
211     return nullptr;
212 
213   // Found an LHS. Parse the remaining expression.
214   auto opLoc = getToken().getLoc();
215   if (AffineHighPrecOp op = consumeIfHighPrecOp()) {
216     if (llhs) {
217       AffineExpr expr = getAffineBinaryOpExpr(llhsOp, llhs, lhs, opLoc);
218       if (!expr)
219         return nullptr;
220       return parseAffineHighPrecOpExpr(expr, op, opLoc);
221     }
222     // No LLHS, get RHS
223     return parseAffineHighPrecOpExpr(lhs, op, opLoc);
224   }
225 
226   // This is the last operand in this expression.
227   if (llhs)
228     return getAffineBinaryOpExpr(llhsOp, llhs, lhs, llhsOpLoc);
229 
230   // No llhs, 'lhs' itself is the expression.
231   return lhs;
232 }
233 
234 /// Parse an affine expression inside parentheses.
235 ///
236 ///   affine-expr ::= `(` affine-expr `)`
237 AffineExpr AffineParser::parseParentheticalExpr() {
238   if (parseToken(Token::l_paren, "expected '('"))
239     return nullptr;
240   if (getToken().is(Token::r_paren))
241     return emitError("no expression inside parentheses"), nullptr;
242 
243   auto expr = parseAffineExpr();
244   if (!expr || parseToken(Token::r_paren, "expected ')'"))
245     return nullptr;
246 
247   return expr;
248 }
249 
250 /// Parse the negation expression.
251 ///
252 ///   affine-expr ::= `-` affine-expr
253 AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
254   if (parseToken(Token::minus, "expected '-'"))
255     return nullptr;
256 
257   AffineExpr operand = parseAffineOperandExpr(lhs);
258   // Since negation has the highest precedence of all ops (including high
259   // precedence ops) but lower than parentheses, we are only going to use
260   // parseAffineOperandExpr instead of parseAffineExpr here.
261   if (!operand)
262     // Extra error message although parseAffineOperandExpr would have
263     // complained. Leads to a better diagnostic.
264     return emitError("missing operand of negation"), nullptr;
265   return (-1) * operand;
266 }
267 
268 /// Returns true if the given token can be represented as an identifier.
269 static bool isIdentifier(const Token &token) {
270   // We include only `inttype` and `bare_identifier` here since they are the
271   // only non-keyword tokens that can be used to represent an identifier.
272   return token.isAny(Token::bare_identifier, Token::inttype) ||
273          token.isKeyword();
274 }
275 
276 /// Parse a bare id that may appear in an affine expression.
277 ///
278 ///   affine-expr ::= bare-id
279 AffineExpr AffineParser::parseBareIdExpr() {
280   if (!isIdentifier(getToken()))
281     return emitWrongTokenError("expected bare identifier"), nullptr;
282 
283   StringRef sRef = getTokenSpelling();
284   for (auto entry : dimsAndSymbols) {
285     if (entry.first == sRef) {
286       consumeToken();
287       return entry.second;
288     }
289   }
290 
291   return emitWrongTokenError("use of undeclared identifier"), nullptr;
292 }
293 
294 /// Parse an SSA id which may appear in an affine expression.
295 AffineExpr AffineParser::parseSSAIdExpr(bool isSymbol) {
296   if (!allowParsingSSAIds)
297     return emitWrongTokenError("unexpected ssa identifier"), nullptr;
298   if (getToken().isNot(Token::percent_identifier))
299     return emitWrongTokenError("expected ssa identifier"), nullptr;
300   auto name = getTokenSpelling();
301   // Check if we already parsed this SSA id.
302   for (auto entry : dimsAndSymbols) {
303     if (entry.first == name) {
304       consumeToken(Token::percent_identifier);
305       return entry.second;
306     }
307   }
308   // Parse the SSA id and add an AffineDim/SymbolExpr to represent it.
309   if (parseElement(isSymbol))
310     return nullptr;
311   auto idExpr = isSymbol
312                     ? getAffineSymbolExpr(numSymbolOperands++, getContext())
313                     : getAffineDimExpr(numDimOperands++, getContext());
314   dimsAndSymbols.push_back({name, idExpr});
315   return idExpr;
316 }
317 
318 AffineExpr AffineParser::parseSymbolSSAIdExpr() {
319   if (parseToken(Token::kw_symbol, "expected symbol keyword") ||
320       parseToken(Token::l_paren, "expected '(' at start of SSA symbol"))
321     return nullptr;
322   AffineExpr symbolExpr = parseSSAIdExpr(/*isSymbol=*/true);
323   if (!symbolExpr)
324     return nullptr;
325   if (parseToken(Token::r_paren, "expected ')' at end of SSA symbol"))
326     return nullptr;
327   return symbolExpr;
328 }
329 
330 /// Parse a positive integral constant appearing in an affine expression.
331 ///
332 ///   affine-expr ::= integer-literal
333 AffineExpr AffineParser::parseIntegerExpr() {
334   auto val = getToken().getUInt64IntegerValue();
335   if (!val.has_value() || (int64_t)val.value() < 0)
336     return emitError("constant too large for index"), nullptr;
337 
338   consumeToken(Token::integer);
339   return builder.getAffineConstantExpr((int64_t)val.value());
340 }
341 
342 /// Parses an expression that can be a valid operand of an affine expression.
343 /// lhs: if non-null, lhs is an affine expression that is the lhs of a binary
344 /// operator, the rhs of which is being parsed. This is used to determine
345 /// whether an error should be emitted for a missing right operand.
346 //  Eg: for an expression without parentheses (like i + j + k + l), each
347 //  of the four identifiers is an operand. For i + j*k + l, j*k is not an
348 //  operand expression, it's an op expression and will be parsed via
349 //  parseAffineHighPrecOpExpression(). However, for i + (j*k) + -l, (j*k) and
350 //  -l are valid operands that will be parsed by this function.
351 AffineExpr AffineParser::parseAffineOperandExpr(AffineExpr lhs) {
352   switch (getToken().getKind()) {
353   case Token::kw_symbol:
354     return parseSymbolSSAIdExpr();
355   case Token::percent_identifier:
356     return parseSSAIdExpr(/*isSymbol=*/false);
357   case Token::integer:
358     return parseIntegerExpr();
359   case Token::l_paren:
360     return parseParentheticalExpr();
361   case Token::minus:
362     return parseNegateExpression(lhs);
363   case Token::kw_ceildiv:
364   case Token::kw_floordiv:
365   case Token::kw_mod:
366     // Try to treat these tokens as identifiers.
367     return parseBareIdExpr();
368   case Token::plus:
369   case Token::star:
370     if (lhs)
371       emitError("missing right operand of binary operator");
372     else
373       emitError("missing left operand of binary operator");
374     return nullptr;
375   default:
376     // If nothing matches, we try to treat this token as an identifier.
377     if (isIdentifier(getToken()))
378       return parseBareIdExpr();
379 
380     if (lhs)
381       emitError("missing right operand of binary operator");
382     else
383       emitError("expected affine expression");
384     return nullptr;
385   }
386 }
387 
388 /// Parse affine expressions that are bare-id's, integer constants,
389 /// parenthetical affine expressions, and affine op expressions that are a
390 /// composition of those.
391 ///
392 /// All binary op's associate from left to right.
393 ///
394 /// {add, sub} have lower precedence than {mul, div, and mod}.
395 ///
396 /// Add, sub'are themselves at the same precedence level. Mul, floordiv,
397 /// ceildiv, and mod are at the same higher precedence level. Negation has
398 /// higher precedence than any binary op.
399 ///
400 /// llhs: the affine expression appearing on the left of the one being parsed.
401 /// This function will return ((llhs llhsOp lhs) op rhs) if llhs is non null,
402 /// and lhs op rhs otherwise; if there is no rhs, llhs llhsOp lhs is returned
403 /// if llhs is non-null; otherwise lhs is returned. This is to deal with left
404 /// associativity.
405 ///
406 /// Eg: when the expression is e1 + e2*e3 + e4, with e1 as llhs, this function
407 /// will return the affine expr equivalent of (e1 + (e2*e3)) + e4, where
408 /// (e2*e3) will be parsed using parseAffineHighPrecOpExpr().
409 AffineExpr AffineParser::parseAffineLowPrecOpExpr(AffineExpr llhs,
410                                                   AffineLowPrecOp llhsOp) {
411   AffineExpr lhs;
412   if (!(lhs = parseAffineOperandExpr(llhs)))
413     return nullptr;
414 
415   // Found an LHS. Deal with the ops.
416   if (AffineLowPrecOp lOp = consumeIfLowPrecOp()) {
417     if (llhs) {
418       AffineExpr sum = getAffineBinaryOpExpr(llhsOp, llhs, lhs);
419       return parseAffineLowPrecOpExpr(sum, lOp);
420     }
421     // No LLHS, get RHS and form the expression.
422     return parseAffineLowPrecOpExpr(lhs, lOp);
423   }
424   auto opLoc = getToken().getLoc();
425   if (AffineHighPrecOp hOp = consumeIfHighPrecOp()) {
426     // We have a higher precedence op here. Get the rhs operand for the llhs
427     // through parseAffineHighPrecOpExpr.
428     AffineExpr highRes = parseAffineHighPrecOpExpr(lhs, hOp, opLoc);
429     if (!highRes)
430       return nullptr;
431 
432     // If llhs is null, the product forms the first operand of the yet to be
433     // found expression. If non-null, the op to associate with llhs is llhsOp.
434     AffineExpr expr =
435         llhs ? getAffineBinaryOpExpr(llhsOp, llhs, highRes) : highRes;
436 
437     // Recurse for subsequent low prec op's after the affine high prec op
438     // expression.
439     if (AffineLowPrecOp nextOp = consumeIfLowPrecOp())
440       return parseAffineLowPrecOpExpr(expr, nextOp);
441     return expr;
442   }
443   // Last operand in the expression list.
444   if (llhs)
445     return getAffineBinaryOpExpr(llhsOp, llhs, lhs);
446   // No llhs, 'lhs' itself is the expression.
447   return lhs;
448 }
449 
450 /// Parse an affine expression.
451 ///  affine-expr ::= `(` affine-expr `)`
452 ///                | `-` affine-expr
453 ///                | affine-expr `+` affine-expr
454 ///                | affine-expr `-` affine-expr
455 ///                | affine-expr `*` affine-expr
456 ///                | affine-expr `floordiv` affine-expr
457 ///                | affine-expr `ceildiv` affine-expr
458 ///                | affine-expr `mod` affine-expr
459 ///                | bare-id
460 ///                | integer-literal
461 ///
462 /// Additional conditions are checked depending on the production. For eg.,
463 /// one of the operands for `*` has to be either constant/symbolic; the second
464 /// operand for floordiv, ceildiv, and mod has to be a positive integer.
465 AffineExpr AffineParser::parseAffineExpr() {
466   return parseAffineLowPrecOpExpr(nullptr, AffineLowPrecOp::LNoOp);
467 }
468 
469 /// Parse a dim or symbol from the lists appearing before the actual
470 /// expressions of the affine map. Update our state to store the
471 /// dimensional/symbolic identifier.
472 ParseResult AffineParser::parseIdentifierDefinition(AffineExpr idExpr) {
473   if (!isIdentifier(getToken()))
474     return emitWrongTokenError("expected bare identifier");
475 
476   auto name = getTokenSpelling();
477   for (auto entry : dimsAndSymbols) {
478     if (entry.first == name)
479       return emitError("redefinition of identifier '" + name + "'");
480   }
481   consumeToken();
482 
483   dimsAndSymbols.push_back({name, idExpr});
484   return success();
485 }
486 
487 /// Parse the list of dimensional identifiers to an affine map.
488 ParseResult AffineParser::parseDimIdList(unsigned &numDims) {
489   auto parseElt = [&]() -> ParseResult {
490     auto dimension = getAffineDimExpr(numDims++, getContext());
491     return parseIdentifierDefinition(dimension);
492   };
493   return parseCommaSeparatedList(Delimiter::Paren, parseElt,
494                                  " in dimensional identifier list");
495 }
496 
497 /// Parse the list of symbolic identifiers to an affine map.
498 ParseResult AffineParser::parseSymbolIdList(unsigned &numSymbols) {
499   auto parseElt = [&]() -> ParseResult {
500     auto symbol = getAffineSymbolExpr(numSymbols++, getContext());
501     return parseIdentifierDefinition(symbol);
502   };
503   return parseCommaSeparatedList(Delimiter::Square, parseElt,
504                                  " in symbol list");
505 }
506 
507 /// Parse the list of symbolic identifiers to an affine map.
508 ParseResult
509 AffineParser::parseDimAndOptionalSymbolIdList(unsigned &numDims,
510                                               unsigned &numSymbols) {
511   if (parseDimIdList(numDims)) {
512     return failure();
513   }
514   if (!getToken().is(Token::l_square)) {
515     numSymbols = 0;
516     return success();
517   }
518   return parseSymbolIdList(numSymbols);
519 }
520 
521 /// Parses an ambiguous affine map or integer set definition inline.
522 ParseResult AffineParser::parseAffineMapOrIntegerSetInline(AffineMap &map,
523                                                            IntegerSet &set) {
524   unsigned numDims = 0, numSymbols = 0;
525 
526   // List of dimensional and optional symbol identifiers.
527   if (parseDimAndOptionalSymbolIdList(numDims, numSymbols))
528     return failure();
529 
530   if (consumeIf(Token::arrow))
531     return parseAffineMapRange(numDims, numSymbols, map);
532 
533   if (parseToken(Token::colon, "expected '->' or ':'"))
534     return failure();
535   return parseIntegerSetConstraints(numDims, numSymbols, set);
536 }
537 
538 /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
539 ParseResult
540 AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
541                                      OpAsmParser::Delimiter delimiter) {
542 
543   SmallVector<AffineExpr, 4> exprs;
544   auto parseElt = [&]() -> ParseResult {
545     auto elt = parseAffineExpr();
546     exprs.push_back(elt);
547     return elt ? success() : failure();
548   };
549 
550   // Parse a multi-dimensional affine expression (a comma-separated list of
551   // 1-d affine expressions); the list can be empty. Grammar:
552   // multi-dim-affine-expr ::= `(` `)`
553   //                         | `(` affine-expr (`,` affine-expr)* `)`
554   if (parseCommaSeparatedList(delimiter, parseElt, " in affine map"))
555     return failure();
556 
557   // Parsed a valid affine map.
558   map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
559                        exprs, getContext());
560   return success();
561 }
562 
563 /// Parse an AffineExpr where the dim and symbol identifiers are SSA ids.
564 ParseResult AffineParser::parseAffineExprOfSSAIds(AffineExpr &expr) {
565   expr = parseAffineExpr();
566   return success(expr != nullptr);
567 }
568 
569 /// Parse the range and sizes affine map definition inline.
570 ///
571 ///  affine-map ::= dim-and-symbol-id-lists `->` multi-dim-affine-expr
572 ///
573 ///  multi-dim-affine-expr ::= `(` `)`
574 ///  multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)`
575 ParseResult AffineParser::parseAffineMapRange(unsigned numDims,
576                                               unsigned numSymbols,
577                                               AffineMap &result) {
578   SmallVector<AffineExpr, 4> exprs;
579   auto parseElt = [&]() -> ParseResult {
580     auto elt = parseAffineExpr();
581     ParseResult res = elt ? success() : failure();
582     exprs.push_back(elt);
583     return res;
584   };
585 
586   // Parse a multi-dimensional affine expression (a comma-separated list of
587   // 1-d affine expressions). Grammar:
588   // multi-dim-affine-expr ::= `(` `)`
589   //                         | `(` affine-expr (`,` affine-expr)* `)`
590   if (parseCommaSeparatedList(Delimiter::Paren, parseElt,
591                               " in affine map range"))
592     return failure();
593 
594   // Parsed a valid affine map.
595   result = AffineMap::get(numDims, numSymbols, exprs, getContext());
596   return success();
597 }
598 
599 /// Parse an affine constraint.
600 ///  affine-constraint ::= affine-expr `>=` `affine-expr`
601 ///                      | affine-expr `<=` `affine-expr`
602 ///                      | affine-expr `==` `affine-expr`
603 ///
604 /// The constraint is normalized to
605 ///  affine-constraint ::= affine-expr `>=` `0`
606 ///                      | affine-expr `==` `0`
607 /// before returning.
608 ///
609 /// isEq is set to true if the parsed constraint is an equality, false if it
610 /// is an inequality (greater than or equal).
611 ///
612 AffineExpr AffineParser::parseAffineConstraint(bool *isEq) {
613   AffineExpr lhsExpr = parseAffineExpr();
614   if (!lhsExpr)
615     return nullptr;
616 
617   // affine-constraint ::= `affine-expr` `>=` `affine-expr`
618   if (consumeIf(Token::greater) && consumeIf(Token::equal)) {
619     AffineExpr rhsExpr = parseAffineExpr();
620     if (!rhsExpr)
621       return nullptr;
622     *isEq = false;
623     return lhsExpr - rhsExpr;
624   }
625 
626   // affine-constraint ::= `affine-expr` `<=` `affine-expr`
627   if (consumeIf(Token::less) && consumeIf(Token::equal)) {
628     AffineExpr rhsExpr = parseAffineExpr();
629     if (!rhsExpr)
630       return nullptr;
631     *isEq = false;
632     return rhsExpr - lhsExpr;
633   }
634 
635   // affine-constraint ::= `affine-expr` `==` `affine-expr`
636   if (consumeIf(Token::equal) && consumeIf(Token::equal)) {
637     AffineExpr rhsExpr = parseAffineExpr();
638     if (!rhsExpr)
639       return nullptr;
640     *isEq = true;
641     return lhsExpr - rhsExpr;
642   }
643 
644   return emitError("expected '== affine-expr' or '>= affine-expr' at end of "
645                    "affine constraint"),
646          nullptr;
647 }
648 
649 /// Parse the constraints that are part of an integer set definition.
650 ///  integer-set-inline
651 ///                ::= dim-and-symbol-id-lists `:`
652 ///                '(' affine-constraint-conjunction? ')'
653 ///  affine-constraint-conjunction ::= affine-constraint (`,`
654 ///                                       affine-constraint)*
655 ///
656 ParseResult AffineParser::parseIntegerSetConstraints(unsigned numDims,
657                                                      unsigned numSymbols,
658                                                      IntegerSet &result) {
659   SmallVector<AffineExpr, 4> constraints;
660   SmallVector<bool, 4> isEqs;
661   auto parseElt = [&]() -> ParseResult {
662     bool isEq;
663     auto elt = parseAffineConstraint(&isEq);
664     ParseResult res = elt ? success() : failure();
665     if (elt) {
666       constraints.push_back(elt);
667       isEqs.push_back(isEq);
668     }
669     return res;
670   };
671 
672   // Parse a list of affine constraints (comma-separated).
673   if (parseCommaSeparatedList(Delimiter::Paren, parseElt,
674                               " in integer set constraint list"))
675     return failure();
676 
677   // If no constraints were parsed, then treat this as a degenerate 'true' case.
678   if (constraints.empty()) {
679     /* 0 == 0 */
680     auto zero = getAffineConstantExpr(0, getContext());
681     result = IntegerSet::get(numDims, numSymbols, zero, true);
682     return success();
683   }
684 
685   // Parsed a valid integer set.
686   result = IntegerSet::get(numDims, numSymbols, constraints, isEqs);
687   return success();
688 }
689 
690 //===----------------------------------------------------------------------===//
691 // Parser
692 //===----------------------------------------------------------------------===//
693 
694 /// Parse an ambiguous reference to either and affine map or an integer set.
695 ParseResult Parser::parseAffineMapOrIntegerSetReference(AffineMap &map,
696                                                         IntegerSet &set) {
697   return AffineParser(state).parseAffineMapOrIntegerSetInline(map, set);
698 }
699 ParseResult Parser::parseAffineMapReference(AffineMap &map) {
700   SMLoc curLoc = getToken().getLoc();
701   IntegerSet set;
702   if (parseAffineMapOrIntegerSetReference(map, set))
703     return failure();
704   if (set)
705     return emitError(curLoc, "expected AffineMap, but got IntegerSet");
706   return success();
707 }
708 ParseResult Parser::parseIntegerSetReference(IntegerSet &set) {
709   SMLoc curLoc = getToken().getLoc();
710   AffineMap map;
711   if (parseAffineMapOrIntegerSetReference(map, set))
712     return failure();
713   if (map)
714     return emitError(curLoc, "expected IntegerSet, but got AffineMap");
715   return success();
716 }
717 
718 /// Parse an AffineMap of SSA ids. The callback 'parseElement' is used to
719 /// parse SSA value uses encountered while parsing affine expressions.
720 ParseResult
721 Parser::parseAffineMapOfSSAIds(AffineMap &map,
722                                function_ref<ParseResult(bool)> parseElement,
723                                OpAsmParser::Delimiter delimiter) {
724   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
725       .parseAffineMapOfSSAIds(map, delimiter);
726 }
727 
728 /// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to parse
729 /// SSA value uses encountered while parsing.
730 ParseResult
731 Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
732                                 function_ref<ParseResult(bool)> parseElement) {
733   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
734       .parseAffineExprOfSSAIds(expr);
735 }
736 
737 IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context,
738                                  bool printDiagnosticInfo) {
739   llvm::SourceMgr sourceMgr;
740   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
741       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
742       /*RequiresNullTerminator=*/false);
743   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
744   SymbolState symbolState;
745   ParserConfig config(context);
746   ParserState state(sourceMgr, config, symbolState, /*asmState=*/nullptr,
747                     /*codeCompleteContext=*/nullptr);
748   Parser parser(state);
749 
750   raw_ostream &os = printDiagnosticInfo ? llvm::errs() : llvm::nulls();
751   SourceMgrDiagnosticHandler handler(sourceMgr, context, os);
752   IntegerSet set;
753   if (parser.parseIntegerSetReference(set))
754     return IntegerSet();
755 
756   Token endTok = parser.getToken();
757   if (endTok.isNot(Token::eof)) {
758     parser.emitError(endTok.getLoc(), "encountered unexpected token");
759     return IntegerSet();
760   }
761 
762   return set;
763 }
764