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
+        sces.append(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))
-        }
       } 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)
     }
   }

Sources/P4Lang/Expressions.swift

@@ -17,91 +17,30 @@
 
 import Common
 
-public class KeysetExpression {
+public struct KeysetExpression {
-  public func update_type(to: P4Type) -> Result<KeysetExpression> {
-    return .Ok(self)
-  }
-
-  public func kse_evaluate(execution: Common.ProgramExecution) -> Result<P4Value> {
-    return .Error(Error(withMessage: "Missing key in keyset expression"))
-  }
-
-  public func kse_type() -> P4Type {
-    return P4Boolean()
-  }
-}
-
-public class NonDefaultKeysetExpression: KeysetExpression {
   public let key: EvaluatableExpression
 
   public init(_ key: EvaluatableExpression) {
     self.key = key
   }
 
-  // Some keyset expressions need additional
+  public func compatible(type: P4Type) -> Result<()> {
-  // context about their types -- e.g., default.
+    if let key_type = self.key.type() as? P4Set {
-  // Override to update and return true if the
+      if !key_type.set_type().eq(rhs: type) {
-  // update is safe.
+        return .Error(
-  public override func update_type(to: P4Type) -> Result<KeysetExpression> {
+          Error(
-    // In the default case, if the current key type
+            withMessage:
-    // does not match the updated type, that's an
+              "Key expression of type set of type \(key_type.set_type()) is not compatible with selector type \(type)"
-    // error.
+          ))
-    return Map(input: key.type().eq(rhs: to)) { input in
+      }
-      input
+    } else if !self.key.type().eq(rhs: type) {
-        ? .Ok(self)
+      return .Error(
-        : .Error(
+        Error(
-          Error(withMessage: "Keyset expression type does not match selector expression type"))
+          withMessage:
+            "Key expression of type \(self.key.type()) is not compatible with selector type \(type)"
+        ))
     }
-  }
+    return .Ok(())
-
-  public override func kse_evaluate(execution: Common.ProgramExecution) -> Result<P4Value> {
-    return self.key.evaluate(execution: execution)
-  }
-
-  public override func kse_type() -> P4Type {
-    return self.key.type()
-  }
-
-}
-
-public class DefaultKeysetExpression: KeysetExpression {
-  let type: P4Type
-
-  public init(withType type: P4Type) {
-    self.type = type
-  }
-
-  public override func update_type(to: P4Type) -> Result<KeysetExpression> {
-    return Map(input: type.eq(rhs: to)) { input in
-      input
-        ? .Ok(DefaultKeysetExpression(withType: to))
-        : .Error(
-          Error(withMessage: "Keyset expression type does not match selector expression type"))
-    }
-  }
-
-  public override func kse_evaluate(execution: Common.ProgramExecution) -> Result<P4Value> {
-    return .Ok(P4SetDefaultValue(withType: self.type))
-  }
-
-  public override func kse_type() -> P4Type {
-    return P4Set(withSetType: self.type)
-  }
-}
-
-public class PlaceholderDefaultKeysetExpression: KeysetExpression {
-  public override init() {}
-
-  public override func update_type(to: P4Type) -> Result<KeysetExpression> {
-    .Ok(DefaultKeysetExpression(withType: to))
-  }
-
-  public override func kse_evaluate(execution: Common.ProgramExecution) -> Result<P4Value> {
-    return .Error(Error(withMessage: "Cannot evaluate a placeholder default keyset expression"))
-  }
-
-  public override func kse_type() -> P4Type {
-    return P4Set(withSetType: P4Boolean())
   }
 }
 
@@ -123,20 +62,6 @@ public struct SelectCaseExpression {
     self.next_state_identifier = next_state_id
     self.next_state = next_state
   }
-
-  // Some keyset expressions need additional
-  // context about their types -- e.g., default.
-  // Override to update and return true if the
-  public func update_type(to: P4Type) -> Result<SelectCaseExpression> {
-    switch key.update_type(to: to) {
-    case .Ok(let new_kse):
-      .Ok(
-        SelectCaseExpression(
-          withKey: new_kse, withNextState: self.next_state_identifier,
-          withNextState: self.next_state))
-    case .Error(let e): .Error(e)
-    }
-  }
 }
 
 public struct SelectExpression {

Sources/P4Runtime/Expressions.swift

@@ -374,12 +374,12 @@ extension FieldAccessExpression: EvaluatableLValueExpression {
 }
 
 extension KeysetExpression: EvaluatableExpression {
-  public func evaluate(execution: Common.ProgramExecution) -> Common.Result<any Common.P4Value> {
+  public func evaluate(execution: Common.ProgramExecution) -> Result<P4Value> {
-    return self.kse_evaluate(execution: execution)
+    return self.key.evaluate(execution: execution)
   }
 
-  public func type() -> any Common.P4Type {
+  public func type() -> P4Type {
-    return self.kse_type()
+    return self.key.type()
   }
 }
 

Tests/p4rseTests/ExpressionTests/FunctionCall.swift

@@ -134,3 +134,24 @@ import TreeSitterP4
 
   #expect(AsInstantiatedParserState(state_result) == P4Lang.reject)
 }
+
+@Test func test_function_call_invalid_return_type() async throws {
+  let simple_parser_declaration = """
+      int functionb(int c) {
+        return true;
+      };
+      parser main_parser() {
+        state start {
+          int c = 5;
+          transition select (4 == functionb(c)) {
+            false: reject;
+            true: accept;
+          };
+        }
+      };
+    """
+
+  let error = try #UseErrorResult(Program.Compile(simple_parser_declaration))
+
+  #expect(error.msg.contains("{29, 12}: Type of expression in return statement (Boolean) is not compatible with function return type (Int)"))
+}

Tests/p4rseTests/ExpressionTests/SelectExpression.swift

@@ -105,7 +105,6 @@ import TreeSitterP4
         }
       };
     """
-
   let program = try #UseOkResult(Program.Compile(simple_parser_declaration))
   let parser = try #UseOkResult(program.find_parser(withName: Identifier(name: "main_parser")))
   let runtime = try #UseOkResult(P4Runtime.ParserRuntime.create(program: program))
@@ -133,7 +132,7 @@ import TreeSitterP4
     #RequireErrorResult(
       Error(
         withMessage:
-        "Error(s) parsing select cases: {81, 12}: Keyset expression type does not match selector expression type"
+        "Error(s) parsing select cases: {81, 4}: Key expression of type Boolean is not compatible with selector type Int"
       ),
       Program.Compile(simple_parser_declaration)))
 }