[Heavy] Finish subpattern expansion

Author: ricejasonf

heavy/include/heavy/Ops.td

@@ -378,23 +378,46 @@ def heavy_MatchPairOp : HeavyOp<"match_pair", []> {
   ];
 }
 
+def heavy_MatchTailOp : HeavyOp<"match_tail", []> {
+  let description = [{
+    Destructure the last N lengthed list of the tail of the input list.
+    The last cdr is counted in the length.
+    (ie () or any element for improper list.)
+    The tail will be matched against a pattern.
+    If the tail is a pair, its address we serve as a sentinel
+    for terminating the matching of elements against the subpattern.
+  }];
+
+  let arguments = (ins UI32Attr:$length,
+                       HeavyValue:$input);
+  let results   = (outs HeavyValue:$tail);
+
+  let builders = [
+    OpBuilder<(ins "uint32_t":$length,
+                   "::mlir::Value":$input)>
+  ];
+}
+
 def heavy_SubpatternOp : HeavyOp<"subpattern", []> {
   let description = [{
     Match a list of elements over a subpattern that
     precedes the `...` ellipsis. Each pattern variable within
     the list is returned as a pack list that may only be expanded
     by a subtemplate.
 
-    The element after the matched list
-    is destructured as `cdr` similar to match_pair.
+    We stop at the elements that occur in the list after the
+    subpattern by passing the tail value as a sentinel.
+    It checks by the address of the pair or any cdr that is
+    not a pair.
   }];
 
-  let arguments = (ins HeavyValue:$input);
-  let results   = (outs Variadic<HeavyValue>:$packs, HeavyValue:$cdr);
+  let arguments = (ins HeavyValue:$input, HeavyValue:$tail);
+  let results   = (outs Variadic<HeavyValue>:$packs);
   let regions   = (region SizedRegion<1>:$body);
 
   let builders = [
     OpBuilder<(ins "::mlir::Value":$input,
+                   "::mlir::Value":$tail,
                    "std::unique_ptr<::mlir::Region>":$body,
                    "unsigned":$num_packs)>
   ];

heavy/lib/Dialect.cpp

@@ -181,14 +181,20 @@ void MatchPairOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
                      input);
 }
 
+void MatchTailOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
+                        uint32_t Length, mlir::Value Input) {
+  mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
+  MatchTailOp::build(B, OpState, HeavyValueT,
+                     Length, Input);
+}
+
 void SubpatternOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
-                        mlir::Value Input,
+                        mlir::Value Input, mlir::Value Tail,
                         std::unique_ptr<mlir::Region> Body,
                         unsigned NumPacks) {
   mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
-  // NumPacks + 1 because $cdr is also a result value.
-  llvm::SmallVector<mlir::Type, 4> ResultTypes(NumPacks + 1, HeavyValueT);
-  OpState.addOperands(Input);
+  llvm::SmallVector<mlir::Type, 4> ResultTypes(NumPacks, HeavyValueT);
+  OpState.addOperands({Input, Tail});
   OpState.addRegion(std::move(Body));
   OpState.addTypes(std::move(ResultTypes));
 }

heavy/lib/OpEval.cpp

@@ -206,6 +206,7 @@ class OpEvalImpl {
     else if (isa<FuncOp>(Op))         return next(Op); // skip functions
     else if (isa<SyntaxOp>(Op))       return Visit(cast<SyntaxOp>(Op));
     else if (isa<MatchPairOp>(Op))    return Visit(cast<MatchPairOp>(Op));
+    else if (isa<MatchTailOp>(Op))    return Visit(cast<MatchTailOp>(Op));
     else if (isa<SubpatternOp>(Op))   return Visit(cast<SubpatternOp>(Op));
     else if (isa<ExpandPacksOp>(Op))  return Visit(cast<ExpandPacksOp>(Op));
     else if (isa<ResolveOp>(Op))      return Visit(cast<ResolveOp>(Op));
@@ -622,7 +623,8 @@ class OpEvalImpl {
   // or, in the case of a pack, it begins matching the rest of the list.
   BlockItrTy patternFail(mlir::Operation* Op) {
     // We should currently be in the scope of a PatternOp or SubpatternOp
-    assert((isa<MatchOp, MatchPairOp, SubpatternOp, MatchIdOp>(Op)) &&
+    assert((isa<MatchOp, MatchPairOp, MatchTailOp,
+                SubpatternOp, MatchIdOp>(Op)) &&
         "Operation must be a pattern matcher");
 
     mlir::Operation* PatternOp = Op->getParentOp();
@@ -672,6 +674,27 @@ class OpEvalImpl {
     return patternFail(Op);
   }
 
+  BlockItrTy Visit(MatchTailOp Op) {
+    // Here we include the last cdr in the length
+    // of a list proper or improper.
+    heavy::Value E = getValue(Op.getInput());
+    uint32_t TargetLen = Op.getLength();
+    assert(TargetLen >= 1 && "expecting positive length");
+    uint32_t TotalLen = 1;
+    heavy::Value Cur = E;
+    while (auto* Pair = dyn_cast<heavy::Pair>(Cur)) {
+      ++TotalLen;
+      Cur = Pair->Cdr;
+    }
+    if (TotalLen < TargetLen)
+      return patternFail(Op);
+    Cur = E;
+    for (uint32_t i = 0; i < TotalLen - TargetLen; i++)
+      Cur = cast<heavy::Pair>(Cur)->Cdr;
+    setValue(Op.getTail(), Cur);
+    return next(Op);
+  }
+
   BlockItrTy Visit(ExpandPacksOp Op) {
     heavy::SourceLocation Loc = getSourceLocation(Op.getLoc());
     setValue(Op.getResult(), getValue(Op.getCdr()));
@@ -708,6 +731,7 @@ class OpEvalImpl {
 
   BlockItrTy Visit(SubpatternOp Op) {
     heavy::Value E = getValue(Op.getInput());
+    heavy::Pair* Tail = dyn_cast<heavy::Pair>(getValue(Op.getTail()));
 
     // Match the empty list.
     if (isa<heavy::Empty>(E)) {
@@ -723,19 +747,20 @@ class OpEvalImpl {
     assert(Op.getBody().getNumArguments() == 1
         && "body should have single argument");
     mlir::Value BodyArg = Op.getBody().getArgument(0);
-    // Visit the subpattern body for each element in E.
+    // Visit the subpattern body for each element in E
+    // stopping when we find tail or a non-pair object.
     // Each "pack" should be a list.
     while (auto* Pair = dyn_cast<heavy::Pair>(E)) {
+      if (Pair == Tail)
+        break;
       push_scope();
       BlockItrTy Itr = Op.getBody().front().begin();
-      setValue(BodyArg, Pair);
+      setValue(BodyArg, Pair->Car);
       while (Itr != BlockItrTy())
         Itr = Visit(&*Itr);
       E = Pair->Cdr;
     }
 
-    setValue(Op.getCdr(), E);
-
     return next(Op);
   }
 
@@ -762,7 +787,7 @@ class OpEvalImpl {
     }
 
     auto ExpandPacksOp = cast<heavy::ExpandPacksOp>(Op->getParentOp());
-    assert(Op.getArgs().size() == 1 && "expecting single result"); 
+    assert(Op.getArgs().size() == 1 && "expecting single result");
     heavy::Value CurrentResult = getValue(Op.getArgs().front());
     pop_scope();
     heavy::Value Result = Context.CreatePair(CurrentResult,

heavy/lib/PatternTemplate.h

@@ -61,7 +61,7 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
   // VisitPatternTemplate should be called with OpGen's insertion point in
   // the body of PatternOp
   mlir::Value VisitPatternTemplate(heavy::Value Pattern,
-                                   heavy::Value Template, 
+                                   heavy::Value Template,
                                    mlir::Value E) {
     heavy::SourceLocation Loc = Pattern.getSourceLocation();
     if (isa_and_nonnull<Symbol>(Pattern.car())) {
@@ -106,8 +106,33 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
     return OpGen.SetError("invalid pattern node", P);
   }
 
+  mlir::Value VisitTail(heavy::SourceLocation Loc,
+                        Value P, mlir::Value E) {
+    // Match the list after the `...`.
+    uint32_t Length = 1;
+
+    if (Pair* P2 = dyn_cast<Pair>(P))
+      Loc = P2->getSourceLocation();
+
+    heavy::Value Cur = P;
+    while (Pair* P2 = dyn_cast<Pair>(Cur)) {
+      ++Length;
+      Cur = P2->Cdr;
+    }
+    mlir::Value Tail = OpGen.create<MatchTailOp>(Loc, Length, E);
+
+    // Actually match the tail part.
+    Visit(P, Tail);
+
+    // The tail is passed to the SubpatternOp.
+    return Tail;
+  }
+
   mlir::Value VisitSubpattern(heavy::SourceLocation Loc,
                               Value P, Value Cdr, mlir::Value E) {
+    // The Tail will be used as a sentinel value if it is a pair.
+    mlir::Value Tail = VisitTail(Loc, Cdr, E);
+
     // Visit the subpattern.
     auto Body = std::make_unique<mlir::Region>();
     llvm::SmallVector<mlir::Value, 4> Packs;
@@ -120,7 +145,7 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
       mlir::Value BodyArg = Block.addArgument(HeavyValueT, MLoc);
       Visit(P, BodyArg);
 
-      // Create range for pack values by finding the 
+      // Create range for pack values by finding the
       // SyntaxClosureOps in the Body region.
       for (mlir::Operation& Op : Block) {
         if (auto SC = dyn_cast<SyntaxClosureOp>(&Op))
@@ -132,11 +157,8 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
     }
 
     // Create the SubpatternOp.
-    auto SubpatternOp = OpGen.create<heavy::SubpatternOp>(
-        Loc, E, std::move(Body), Packs.size());
-
-    if (!OpGen.CheckError())
-      Visit(Cdr, SubpatternOp.getCdr());
+    OpGen.create<heavy::SubpatternOp>(
+        Loc, E, Tail, std::move(Body), Packs.size());
 
     // In the template, the nested syntax closures
     // will be looked up and check if its parent is
@@ -169,7 +191,7 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
     if (P->equals("_")) {
       // Since _ always matches anything, there is
       // nothing to check.
-      return mlir::Value(); 
+      return mlir::Value();
     }
 
     // <pattern identifier> (literal identifier)

heavy/lib/TemplateGen.h

@@ -49,7 +49,7 @@ class TemplateGen : TemplateBase<TemplateGen>,
   void VisitTemplate(heavy::Value Template) {
     heavy::SourceLocation Loc = Template.getSourceLocation();
     ResultTy Result = Visit(Template);
-    
+
     mlir::Value TransformedSyntax;
     if (mlir::Value* MVP = std::get_if<mlir::Value>(&Result))
       TransformedSyntax = *MVP;

heavy/test/Evaluate/define-syntax.scm

@@ -1,6 +1,8 @@
 ; RUN: heavy-scheme %s | FileCheck %s
 (import (heavy base))
 
+;; TODO Local let-syntax needs to be tested/supported.
+
 (define ok 'ok!)
 (define-syntax my-lambda
   (syntax-rules (=>)
@@ -17,16 +19,38 @@
  'oops!)
 (newline)
 
-;; TODO Local let-syntax needs to be supported.
-
-#|
+; CHECK: 42"x has type Int""y has type Int"
+; CHECK-NEXT: 5
+; CHECK-NEXT: 6
 (define-syntax my-lambda
   (syntax-rules (:)
-    ((my-lambda (arg ...) body ...)
+    ((my-lambda ((arg : type) ...) body ...)
      (lambda (arg ...)
+       (write 42)
        (write (string-append 'arg " has type " 'type)) ...
+       (newline)
        body ...))))
-((my-lambda ((x) (y))
+((my-lambda ((x : Int) (y : Int))
   (write x)
-  (write y)))
-|#
+  (newline)
+  (write y)) 5 6)
+(newline)
+
+; CHECK: (0 1)
+; CHECK-NEXT: (0 1 2)
+(define-syntax ez
+  (syntax-rules ()
+    ((ez 0 1 i ...)
+     '(0 1 i ...))))
+(write (ez 0 1))
+(newline)
+(write (ez 0 1 2))
+(newline)
+
+; CHECK: (0 1 2 3 4 9)
+(define-syntax ez
+  (syntax-rules ()
+    ((ez 0 1 i ... 5 6)
+     '(0 1 i ... 9))))
+(write (ez 0 1 2 3 4 5 6))
+(newline)