gp4/Sources/P4Parser/Parsers/Statement.swift

// p4rse, Copyright 2026, Will Hawkins
//
// This file is part of p4rse.
//
// This file is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

import Common
import SwiftTreeSitter
import TreeSitterExtensions
import TreeSitterP4

extension CST.BlockStatement: Parsable {
  public typealias C = CST.BlockStatement
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.BlockStatement> {
    /*
    #RequireNodeType<Node, AST.BlockStatement>(
      node: node, type: "blockStatement", nice_type_name: "block statement")

    var walker = Walker(node: node)
    var current_node: Node? = .none

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<AST.BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType != "{" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: current_node!.toSourceLocation(),
          withError: "Missing { on block statement"))
    }

    var statements: [AST.AnStatement] = Array()
    var parse_err: (any Errorable)? = .none

    walker.next()
    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<AST.BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType == "statements" {
      switch SpecialCompilers.Statements.Parse(
        node: current_node!, withContext: context)
      {
      case .Ok(let parsed_statements):
        statements = parsed_statements
      case .Error(let error):
        parse_err = error
      }

      walker.next()
    }

    if let err = parse_err {
      return .Error(err)
    }

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<AST.BlockStatement>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Malformed block statement")))

    if current_node!.nodeType != "}" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: current_node!.toSourceLocation(),
          withError: "Missing } on block statement"))
    }

    return .Ok(AST.BlockStatement(statements))
    */
    return .Ok(CST.BlockStatement(CST.Statements([])))
  }
}

@deriveParsableStatement
extension CST.BlockStatement: ParsableStatement {}

extension CST.ConditionalStatement: Parsable {
  public typealias C = CST.ConditionalStatement
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.ConditionalStatement> {

    #RequireNodeType<Node, CST.ConditionalStatement>(
      node: node, type: "conditionalStatement", nice_type_name: "conditional statement")

    let maybe_condition_expression = node.child(at: 2)
    guard let condition_expression = maybe_condition_expression,
      condition_expression.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find condition for conditional statement"))
    }

    let maybe_thens = node.child(at: 4)
    guard let thens = maybe_thens,
      thens.nodeType == "statement"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find then statement block for conditional statement"))
    }

    guard
      case .Ok(let condition) = CST.Expression.Parse(
        node: condition_expression, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse a conditional expression in a conditional statement"))
    }

    guard
      case .Ok(let thenns) = CST.Statement.Parse(
        node: thens, withContext: context)
    else {
      return Result.Error(
        Error(
          withMessage:
            "Could not parse the then block in a conditional statement"))
    }

    let optional_elss: Result<CST.Categories.Statement>? =
      if let elss = node.child(at: 6) {
        .some(
          CST.Statement.Parse(
            node: elss, withContext: context))
      } else {
        .none
      }

    if let parsed_elss = optional_elss {
      guard
        case .Ok(let elss) = parsed_elss
      else {
        return Result.Error(
          Error(
            withMessage:
              "Could not parse the else block in a conditional statement"))
      }
      return .Ok(
        CST.ConditionalStatement(condition: condition, withThen: thenns, andElse: elss))
    }
    return .Ok(CST.ConditionalStatement(condition: condition, withThen: thenns))
  }
}

@deriveParsableStatement
extension CST.ConditionalStatement: ParsableStatement {}

extension CST.VariableDeclarationStatement: Parsable {
  public typealias C = CST.VariableDeclarationStatement
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.VariableDeclarationStatement> {

    #RequireNodeType<Node, CST.VariableDeclarationStatement>(
      node: node, type: "variableDeclaration", nice_type_name: "variable declaration statement")

    let maybe_typeref = node.child(at: 0)
    guard let typeref = maybe_typeref,
      typeref.nodeType == "typeRef"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find type name for variable declaration statement"))
    }

    let maybe_variablename = node.child(at: 1)
    guard let variablename = maybe_variablename,
      variablename.nodeType == "identifier"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Did not find identifier name for variable declaration statement"))
    }

    let maybe_rvalue = node.childCount > 3 ? node.child(at: 3) : .none

    guard
      case .Ok(let parsed_variablename) = CST.Identifier.ParseExpression(
        node: variablename, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse variable name"))
    }

    guard
      case .Ok(let declaration_p4_type) = CST.Types.ParseType(type: typeref, withContext: context)
    else {
      return Result.Error(
        Error(withMessage: "Could not parse a P4 type from \(typeref.text!)"))
    }

    var initializer: CST.Categories.Expression? = .none

    // If there is an initializer, it must be an expression.
    if let initializer_expression = maybe_rvalue {
      guard initializer_expression.nodeType == "expression" else {
        return Result.Error(
          ErrorWithLocation(
            sourceLocation: node.toSourceLocation(),
            withError: "initial value for declaration statement is not an expression"))
      }

      let maybe_parsed_rvalue = CST.Expression.Parse(
        node: initializer_expression, withContext: context)
      guard
        case .Ok(let parsed_rvalue) = maybe_parsed_rvalue
      else {
        return .Error(maybe_parsed_rvalue.error()!)
      }

      initializer = parsed_rvalue
    }

    return Result.Ok(
      CST.VariableDeclarationStatement(
        identifier: parsed_variablename as! CST.Identifier, withType: declaration_p4_type,
        withInitializer: initializer),
    )
  }
}

@deriveParsableStatement
extension CST.VariableDeclarationStatement: ParsableStatement {}

extension CST.ExpressionStatement: Parsable {
  public typealias C = CST.ExpressionStatement
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.ExpressionStatement> {
    #RequireNodeType<Node, (P4Statement)>(
      node: node, type: "expressionStatement", nice_type_name: "expression statement")

    let expression_node = node.child(at: 0)!

    return switch CST.Expression.Parse(node: expression_node, withContext: context) {
    case .Ok(let expression): .Ok(CST.ExpressionStatement(expression))
    case .Error(let e): .Error(e)
    }
  }
}

@deriveParsableStatement
extension CST.ExpressionStatement: ParsableStatement {}

extension CST.ParserAssignmentStatement: Parsable {
  public typealias C = CST.ParserAssignmentStatement
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.ParserAssignmentStatement> {

    #RequireNodeType<Node, CST.ParserAssignmentStatement>(
      node: node, type: "assignmentStatement", nice_type_name: "assignment statement")

    guard let lvalue_node = node.child(at: 0),
      lvalue_node.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing lvalue in assignment statement"))
    }

    guard let rvalue_node = node.child(at: 2),
      rvalue_node.nodeType == "expression"
    else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Missing rvalue in assignment statement"))
    }

    let maybe_parsed_rvalue = CST.Expression.Parse(
      node: rvalue_node, withContext: context)
    guard case Result.Ok(let rvalue) = maybe_parsed_rvalue else {
      return Result.Error(maybe_parsed_rvalue.error()!)
    }

    let maybe_parsed_lvalue = CST.Expression.Parse(node: lvalue_node, withContext: context)
    guard case .Ok(let lvalue) = maybe_parsed_lvalue else {
      return Result.Error(maybe_parsed_lvalue.error()!)
    }

    return Result.Ok(
      CST.ParserAssignmentStatement(
        withLValue: lvalue,
        withValue: rvalue
      ))
  }
}

@deriveParsableStatement
extension CST.ParserAssignmentStatement: ParsableStatement {}

extension CST.ReturnStatement: Parsable {
  public typealias C = CST.ReturnStatement
  public static func Parse(
    node: SwiftTreeSitter.Node, withContext context: CSTCompilerContext
  ) -> Result<CST.ReturnStatement> {
    #RequireNodeType<Node, CST.ReturnStatement>(
      node: node, type: "return_statement", nice_type_name: "return statement")

    let expression_node = node.child(at: 1)!

    return switch CST.Expression.Parse(node: expression_node, withContext: context) {
    case .Ok(let result): .Ok(CST.ReturnStatement(result))
    case .Error(let e): .Error(e)
    }
  }
}

@deriveParsableStatement
extension CST.ReturnStatement: ParsableStatement {}

extension CST.ApplyStatement: Parsable {
  public typealias C = CST.ApplyStatement
  public static func Parse(
    node: SwiftTreeSitter.Node, withContext context: CSTCompilerContext
  ) -> Result<CST.ApplyStatement> {
    #RequireNodeType<Node, CST.ApplyStatement>(
      node: node, type: "apply_statement", nice_type_name: "apply statement")

    let expression_node = node.child(at: 1)!

    return switch CST.BlockStatement.Parse(node: expression_node, withContext: context) {
    case .Ok(let statement):
      .Ok(CST.ApplyStatement(statement))
    case .Error(let e): .Error(e)
    }
  }
}

@deriveParsableStatement
extension CST.ApplyStatement: ParsableStatement {}

extension CST.Instantiation: Parsable {
  public typealias C = CST.Instantiation
  public static func Parse(
    node: SwiftTreeSitter.Node, withContext context: CSTCompilerContext
  ) -> Result<CST.Instantiation> {

    let expression = node
    #RequireNodeType<Node, CST.Instantiation>(
      node: expression, type: "instantiation", nice_type_name: "instantiation statement")

    var walker = Walker(node: expression)
    var current_node: Node? = .none

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<CST.Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing function call component")))

    let maybe_instantiated_type_name = CST.Identifier.ParseExpression(
      node: current_node!, withContext: context)
    guard case .Ok(let instantiated_type_name) = maybe_instantiated_type_name else {
      return Result.Error(maybe_instantiated_type_name.error()!)
    }
    walker.next()

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<CST.Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing instantiation component")))

    let maybe_argument_list = CST.ArgumentList.Parse(node: current_node!, withContext: context)

    guard case .Ok(let arguments) = maybe_argument_list else {
      return .Error(maybe_argument_list.error()!)
    }

    walker.next()

    #MustOr(
      result: current_node, thing: walker.getNext(),
      or: Result<CST.Instantiation>.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing instantiation name")))

    let name = CST.Identifier.ParseExpression(node: current_node!, withContext: context)
    guard case .Ok(let name) = name else {
      return .Error(name.error()!)
    }

    return .Ok(
      CST.Instantiation(
        named: name as! CST.Identifier, withType: instantiated_type_name as! CST.Identifier,
        withArguments: arguments))
  }
}

@deriveParsableStatement
extension CST.Instantiation: ParsableStatement {}

extension CST.Statements: Parsable {
  public typealias C = CST.Statements
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.Statements> {
    if node.nodeType != "statements" && node.nodeType != "parserStatements" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Did not find expected statements"))
    }

    var errors: (any Errorable)? = .none
    var parsed_s: [CST.Categories.Statement] = Array()

    node.enumerateNamedChildren { node in
      switch CST.Statement.Parse(
        node: node, withContext: context)
      {
      case .Ok(let parsed_statement):
        parsed_s.append(parsed_statement)
      case .Error(let e):
        errors =
          if let errors = errors {
            errors.append(error: e)
          } else {
            e
          }
      }
    }

    if let errors = errors {
      return .Error(errors)
    }

    return Result.Ok(CST.Statements(parsed_s))
  }
}

extension CST.TransitionStatement: Parsable {
  public static func Parse(
    node: Node, withContext context: CSTCompilerContext
  ) -> Result<CST.Categories.State> {

    guard let state_identifier = context.lexical_context_name else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Cannot parse a transition statement without the name of the containing state."
        ))
    }

    let stmts = context.lexical_context_statements

    #RequireNodeType<Node, P4Statement>(
      node: node, type: "parserTransitionStatement", nice_type_name: "parser transition statement"
    )

    guard let tse_node = node.child(at: 1),
      tse_node.nodeType! == "transitionSelectionExpression"
    else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Could not find transition select expression"))
    }

    guard let next_node = tse_node.child(at: 0) else {
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Could not find the next token in a transition selection expression"))
    }

    // If the next node is an identifier, we have the simple form ...
    if next_node.nodeType == "identifier" {
      let maybe_parsed_next_state_id = CST.Identifier.ParseExpression(
        node: next_node, withContext: context)
      switch maybe_parsed_next_state_id {
      case .Ok(let next_state_id):
        return .Ok(
          CST.ParserStateDirectTransition(
            name: (state_identifier),
            withNextStateIdentifier: next_state_id as! CST.Identifier, withStatements: stmts))
      case .Error(let e):
        return .Error(e)
      }
    }

    // We know that the next node is a select expression.
    return
      switch CST.SelectExpression.ParseExpression(node: next_node, withContext: context)
    {
    case .Ok(let tse):
      .Ok(
        CST.ParserStateSelectTransition(
          name: state_identifier, withTransitionExpression: tse as! CST.SelectExpression,
          withStatements: stmts,
        )
      )
    case .Error(let e): .Error(e)
    }
  }
}

extension CST.Statement: Parsable {
  public typealias C = CST.Categories.Statement
  public static func Parse(
    node: SwiftTreeSitter.Node, withContext context: CSTCompilerContext
  ) -> Result<CST.Categories.Statement> {

    if node.nodeType != "parserStatement" && node.nodeType != "statement" {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(), withError: "Missing expected parser statement")
      )
    }

    let statement = node.child(at: 0)!

    let statementParsers: [String: ParsableStatement.Type] = [
      "assignmentStatement": CST.ParserAssignmentStatement.self,
      "expressionStatement": CST.ExpressionStatement.self,
      "variableDeclaration": CST.VariableDeclarationStatement.self,
      "conditionalStatement": CST.ConditionalStatement.self,
      "blockStatement": CST.BlockStatement.self,
      "return_statement": CST.ReturnStatement.self,
    ]
    guard let parser = statementParsers[statement.nodeType ?? ""] else {
      return Result.Error(
        ErrorWithLocation(
          sourceLocation: statement.toSourceLocation(),
          withError:
            "Unparseable statement type (\(statement.nodeType ?? "Unknown Statement Type"))"))
    }
    switch parser.ParseStatement(node: statement, withContext: context) {
    case Result.Ok(let parsed):
      return .Ok(parsed)
    case Result.Error(let e):
      return .Error(
        ErrorWithLocation(
          sourceLocation: node.toSourceLocation(),
          withError: "Failed to parse a statement element: \(e)"))
    }
  }
}