Refactor Expected Types During Compilation

By adding an expected type to the compilation context, it is
now possible for type checking to occur on keyset expressions
and return statements at the moment that they are being compiled.

Previously, it was necessary to tentatively compile them and then
typecheck afterward.

Signed-off-by: Will Hawkins <hawkinsw@obs.cr>

Sources/P4Compiler/Compiler.swift

@@ -40,39 +40,72 @@ public func ErrorOnNode(node: Node, withError error: String) -> Error {
   return Error(withMessage: "\(node.range): \(error)")
 }
 
-/// Context for compilation.
+/// Context for compilation
+///
+/// It contains (at least) three important pieces of information:
+/// 1. Instances: A ``VarTypeScopes`` that contains information about instantiated objects
+/// (and their types) in scope
+/// 1. Types: A ``TypeTypeScopes`` that contains information about declared types in scope.
+/// 1. Expected Type: In certain situations, to typecheck an element of a P4 program requires
+/// knowledge of an expected type. For instance, when compiling a return statement, the
+/// compiler must know the return type of the function to type check.
 public struct CompilerContext {
   let instances: VarTypeScopes
   let types: TypeTypeScopes
+  let expected_type: P4Type?
 
   public init(withInstances _instances: VarTypeScopes) {
     instances = _instances
     types = TypeTypeScopes()
+    expected_type = .none
   }
 
   public init(withInstances _instances: VarTypeScopes, withTypes _types: TypeTypeScopes) {
     instances = _instances
     types = _types
+    expected_type = .none
+  }
+
+  public init(
+    withInstances _instances: VarTypeScopes, withTypes _types: TypeTypeScopes,
+    withExpectation expectation: P4Type?
+  ) {
+    instances = _instances
+    types = _types
+    expected_type = expectation
   }
 
   /// Update a compiler context
   ///
   /// Create a new compiler context based on the current with the same types and new names.
   ///
-  /// - Parameter names: a ``TypeScopes`` with the updated names for the newly created compiler context.
+  /// - Parameter instances: a ``VarTypeScopes`` with the updated instances for the newly created compiler context.
   /// - Returns: A new compiler context based on the current with the same types and new names.
   public func update(newInstances instances: VarTypeScopes) -> CompilerContext {
-    return CompilerContext(withInstances: instances, withTypes: self.types)
+    return CompilerContext(
+      withInstances: instances, withTypes: self.types, withExpectation: self.expected_type)
   }
 
   /// Update a compiler context
   ///
   /// Create a new compiler context based on the current with the same names and new types.
   ///
-  /// - Parameter types: a ``TypeScopes`` with the updated types for the newly created compiler context.
+  /// - Parameter types: a ``TypeTypeScopes`` with the updated types for the newly created compiler context.
   /// - Returns: A new compiler context based on the current with the same names and new types.
   public func update(newTypes types: TypeTypeScopes) -> CompilerContext {
-    return CompilerContext(withInstances: self.instances, withTypes: types)
+    return CompilerContext(
+      withInstances: self.instances, withTypes: types, withExpectation: self.expected_type)
+  }
+
+  /// Update a compiler context
+  ///
+  /// Create a new compiler context based on the current with the same names and types but new expected type.
+  ///
+  /// - Parameter expectation: a ``P4Type?`` to (re)set the type the compiler is expecting.
+  /// - Returns: A new compiler context based on the current with the same names and types but new expected type.
+  public func update(newExpectation expectation: P4Type?) -> CompilerContext {
+    return CompilerContext(
+      withInstances: self.instances, withTypes: self.types, withExpectation: expectation)
   }
 
 }

Sources/P4Compiler/Declarations.swift

@@ -112,7 +112,9 @@ extension FunctionDeclaration: CompilableDeclaration {
     }
 
     let maybe_function_body = Parser.Statement.Compile(
-      node: currentChild!, withContext: context.update(newInstances: function_scope))
+      node: currentChild!,
+      withContext: context.update(newInstances: function_scope).update(
+        newExpectation: function_type))
     guard case .Ok((let function_body, _)) = maybe_function_body else {
       return .Error(maybe_function_body.error()!)
     }

Sources/P4Compiler/Expression.swift

@@ -129,7 +129,7 @@ extension KeysetExpression: CompilableExpression {
 
     // If there is a default keyset, that's easy!
     if keyset_expression_node.nodeType == "default_keyset" {
-      return .Ok(PlaceholderDefaultKeysetExpression())
+      return .Ok(KeysetExpression(P4SetDefaultValue(withType: context.expected_type!)))
     }
 
     // Compile the expression:
@@ -139,7 +139,7 @@ extension KeysetExpression: CompilableExpression {
       return .Error(maybe_compiled_set_expression.error()!)
     }
 
-    return .Ok(NonDefaultKeysetExpression(compiled_expression))
+    return .Ok(KeysetExpression(compiled_expression))
   }
 }
 
@@ -261,13 +261,9 @@ extension SelectExpression: CompilableExpression {
 
     select_body_node.enumerateNamedChildren { current_node in
       let maybe_parsed_cse = SelectCaseExpression.compile(
-        node: current_node, withContext: context)
+        node: current_node, withContext: context.update(newExpectation: selector.type()))
       if case .Ok(let parsed_cse) = maybe_parsed_cse {
-        let parsed_cse = parsed_cse as! SelectCaseExpression
-        switch parsed_cse.update_type(to: selector.type()) {
-        case .Ok(let updated_cse): sces.append(updated_cse)
-        case .Error(let e): sces_errors.append(ErrorOnNode(node: current_node, withError: e.msg))
-        }
+        sces.append(parsed_cse as! SelectCaseExpression)
       } else {
         sces_errors.append(Error(withMessage: "\(maybe_parsed_cse.error()!)"))
       }
@@ -309,10 +305,21 @@ extension SelectCaseExpression: CompilableExpression {
 
     let maybe_parsed_keysetexpression = KeysetExpression.compile(
       node: keysetexpression_node, withContext: context)
-    guard case Result.Ok(let keysetexpression) = maybe_parsed_keysetexpression else {
+    guard case Result.Ok(let maybe_keysetexpression) = maybe_parsed_keysetexpression else {
       return Result.Error(maybe_parsed_keysetexpression.error()!)
     }
 
+    guard let maybe_keysetexpression = maybe_keysetexpression else {
+      return Result.Error(
+        ErrorOnNode(node: keysetexpression_node, withError: "Missing expected keyset expression"))
+    }
+
+    let keysetexpression = maybe_keysetexpression as! KeysetExpression
+
+    if case .Error(let e) = keysetexpression.compatible(type: context.expected_type!) {
+      return .Error(ErrorOnNode(node: keysetexpression_node, withError: e.msg))
+    }
+
     let maybe_parsed_targetstate = Identifier.Compile(
       node: targetstate_node, withContext: context)
     guard case .Ok(let targetstate) = maybe_parsed_targetstate else {
@@ -321,7 +328,7 @@ extension SelectCaseExpression: CompilableExpression {
 
     return .Ok(
       SelectCaseExpression(
-        withKey: keysetexpression as! KeysetExpression, withNextState: targetstate)
+        withKey: keysetexpression, withNextState: targetstate)
     )
   }
 }

Sources/P4Compiler/Statement.swift

@@ -310,7 +310,17 @@ extension ReturnStatement: CompilableStatement {
     let expression_node = node.child(at: 1)!
 
     return switch Expression.Compile(node: expression_node, withContext: context) {
-    case .Ok(let result): .Ok((ReturnStatement(result), context))
+    case .Ok(let result):
+      if result.type().eq(rhs: context.expected_type!) {
+        .Ok((ReturnStatement(result), context))
+      } else {
+        .Error(
+          ErrorOnNode(
+            node: node,
+            withError:
+              "Type of expression in return statement (\(result.type())) is not compatible with function return type (\(context.expected_type!))"
+          ))
+      }
     case .Error(let e): .Error(e)
     }
   }