[Heavy] Add type checking; Check for procedures before function application

Author: ricejasonf

heavy/include/heavy/Dialect.h

@@ -125,6 +125,7 @@ Heavy##NAME##Ty, mlir::Type, mlir::TypeStorage> { \
 } \
 
 HEAVY_TYPE(Pair, "heavy.pair", "pair");
+HEAVY_TYPE(Procedure, "heavy.procedure", "procedure");
 
 HEAVY_TYPE(Syntax, "heavy.syntax", "syntax");
 HEAVY_TYPE(OpGen, "heavy.opgen", "opgen");

heavy/include/heavy/OpGen.h

@@ -184,6 +184,9 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
 
   mlir::ModuleOp getModuleOp();
+  
+  mlir::Value CheckType(heavy::SourceLocation Loc, mlir::Value V,
+                        mlir::Type Type);
 
   // GetSingleResult
   //  - visits a node expecting a single result

heavy/include/heavy/Ops.td

@@ -45,6 +45,11 @@ def HeavyValueRefs : HeavyValueBase<"ValueRefs">; // For arguments array.
 def HeavyRest : HeavyValueBase<"Rest">; // For rest arguments.
 def HeavyPair : HeavyValueBase<"Pair">;
 def HeavySyntax : HeavyValueBase<"Syntax">;
+def HeavyProcedure : HeavyValueBase<"Procedure">;
+
+def HeavyType : AnyTypeOf<[HeavyValue,
+                           HeavyPair,
+                           HeavyProcedure]>;
 
 // If !heavy.value_refs is a function argument, the arguments should
 // only be (!heavy.context, !heavy.value_refs).
@@ -67,7 +72,7 @@ def heavy_ApplyOp : HeavyOp<"apply", [Terminator]> {
     of the current call arguments with arguments in that position.
   }];
 
-  let arguments = (ins HeavyValue:$fn,
+  let arguments = (ins AnyTypeOf<[HeavyValue, HeavyProcedure]>:$fn,
     Variadic<AnyTypeOf<[HeavyValue, HeavyRest, HeavyValueRefs]>>:$args);
   let results   = (outs);
 }
@@ -437,6 +442,17 @@ def heavy_MatchArgsOp : HeavyOp<"match_args"> {
   ];
 }
 
+def heavy_MatchTypeOp : HeavyOp<"match_type", []> {
+  let description = [{
+    Perform dynamic cast resulting in a value with a specified type.
+    Like other Match operations this will raise an error unless it
+    appears in the body of a PatternOp.
+  }];
+
+  let arguments = (ins HeavyValue:$arg);
+  let results   = (outs HeavyType:$result);
+}
+
 def heavy_SubpatternOp : HeavyOp<"subpattern", []> {
   let description = [{
     Match a list of elements over a subpattern that

heavy/lib/Dialect.cpp

@@ -31,6 +31,7 @@ Dialect::Dialect(mlir::MLIRContext* Ctx)
   addTypes<HeavyValueRefsTy>();
   addTypes<HeavyRestTy>();
   addTypes<HeavyPairTy>();
+  addTypes<HeavyProcedureTy>();
 
   addTypes<HeavySyntaxTy>();
   addTypes<HeavyOpGenTy>();
@@ -97,6 +98,8 @@ void Dialect::printType(mlir::Type Type,
     Name = "rest";
   } else if (mlir::isa<HeavyPairTy>(Type)) {
     Name = "pair";
+  } else if (mlir::isa<HeavyProcedureTy>(Type)) {
+    Name = "procedure";
   } else if (mlir::isa<HeavySyntaxTy>(Type)) {
     Name = "syntax";
   } else if (mlir::isa<HeavyOpGenTy>(Type)) {

heavy/lib/OpEval.cpp

@@ -198,6 +198,7 @@ class OpEvalImpl {
     else if (isa<MatchPairOp>(Op))    return Visit(cast<MatchPairOp>(Op));
     else if (isa<MatchTailOp>(Op))    return Visit(cast<MatchTailOp>(Op));
     else if (isa<MatchArgsOp>(Op))    return Visit(cast<MatchArgsOp>(Op));
+    else if (isa<MatchTypeOp>(Op))    return Visit(cast<MatchTypeOp>(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));
@@ -656,7 +657,7 @@ class OpEvalImpl {
   BlockItrTy patternFail(mlir::Operation* Op, llvm::StringRef ErrMsg,
                          llvm::ArrayRef<heavy::Value> Irr) {
     assert((isa<MatchOp, MatchPairOp, MatchTailOp, MatchArgsOp,
-                MatchIdOp, SubpatternOp>(Op)) &&
+                MatchTypeOp, MatchIdOp, SubpatternOp>(Op)) &&
         "Operation must be a pattern matcher");
 
     mlir::Operation* ParentOp = Op->getParentOp();
@@ -775,6 +776,34 @@ class OpEvalImpl {
       return next(Op);
   }
 
+  BlockItrTy Visit(MatchTypeOp Op) {
+    mlir::Type Type = Op.getResult().getType();
+    heavy::Value Arg = getValue(Op.getArg());
+    bool Result;
+    switch (Arg.getKind()) {
+      case ValueKind::Lambda:
+      case ValueKind::Builtin:
+        Result = isa<HeavyProcedureTy>(Type);
+        break;
+      case ValueKind::Pair:
+        Result = isa<HeavyPairTy>(Type);
+        break;
+      default:
+        Result = false;
+    }
+
+    if (!Result) {
+      std::string Str;
+      llvm::raw_string_ostream Stream(Str);
+      Type.print(Stream);
+      Symbol* TypeStr = Context.CreateSymbol(Str);
+      return patternFail(Op, "expecting object with type {}: {}",
+                         {TypeStr, Arg});
+    }
+    setValue(Op.getResult(), Arg);
+    return next(Op);
+  }
+
   BlockItrTy Visit(ExpandPacksOp Op) {
     heavy::SourceLocation Loc = getSourceLocation(Op.getLoc());
     setValue(Op.getResult(), getValue(Op.getCdr()));

heavy/lib/OpGen.cpp

@@ -185,6 +185,12 @@ mlir::Value OpGen::GetSingleResult(heavy::Value V) {
   return LocalizeValue(Result);
 }
 
+// Insert a check to assert the type of a value.
+mlir::Value OpGen::CheckType(heavy::SourceLocation Loc, mlir::Value V,
+                             mlir::Type Type) {
+  return create<MatchTypeOp>(Loc, Type, V);
+}
+
 // WithLibraryEnv - Call a thunk within the library environment
 //                  for <library spec>. (ie begin, import, export)
 void OpGen::WithLibraryEnv(Value Thunk) {
@@ -1259,6 +1265,9 @@ mlir::Value OpGen::HandleCall(Pair* P, heavy::EnvEntry FnEnvEntry) {
     else
       Fn = GetSingleResult(P->Car);
 
+    // Insert the check that Fn is a procedure.
+    Fn = CheckType(Loc, Fn, Builder.getType<HeavyProcedureTy>());
+
     if (CheckError())
       return Error();