gp4/Sources/Macros/Macros.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 SwiftCompilerPlugin
import SwiftSyntax
@_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros

public func remove_embedded_quotes(_ from: String) -> String {
  let result = from.replacing("\"", with: [])
  return result
}

struct MacroError: Error, CustomStringConvertible {
  var message: String
  var description: String {
    return message
  }
  public init(withMessage _message: String) {
    message = _message
  }
}

public struct UseOkResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Ok(let __good): return __good
              case Result.Error(let __error):
                  print("Unexpected result: \\(__error)")
                  throw Require.Error.UnexpectedResult
          }
      }()
      """
  }
}

public struct UseErrorResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Error(let __error): return __error
              case Result.Ok(let __good):
                  print("Unexpected result: \\(__good)")
                  throw Require.Error.UnexpectedResult
          }
      }()
      """
  }
}

public struct Require {
  public enum Error: Swift.Error {
    case UnexpectedResult
    case SyntaxError
  }
}

public struct RequireResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    guard let argument = node.arguments.first?.expression else {
      throw Require.Error.SyntaxError
    }

    return """
      {
          switch \(argument) {
              case Result.Ok(_): return true
              case Result.Error(let __error):
                  print("Unexpected result: \\(__error)")
                  return false
          }
      }()
      """
  }
}

public struct RequireErrorResult: ExpressionMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax,
    in context: some MacroExpansionContext
  ) throws -> ExprSyntax {

    let arguments = node.arguments.indices
    let expected_error = node.arguments[arguments.startIndex].expression
    let error_producer = node.arguments[arguments.index(after: arguments.startIndex)].expression

    return ExprSyntax(
      """
      {
          let __expected_error = \(expected_error)
          let __actual_error = \(error_producer)
          if case Result.Error(__expected_error) = __actual_error {
              return true
          } else {
              print("Expected Error: \\(__expected_error) but got Error: \\(__actual_error)")
              return false
          }
      }()
      """)
  }
}

public struct RequireNodeType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex
    let node_to_check = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type_nice_name = node.arguments[arg_index].expression

    let error_message =
      "Did not find " + remove_embedded_quotes(expected_type_nice_name.description)

    return [
      CodeBlockItemSyntax(
        """
        if \(node_to_check).nodeType != \(expected_type) {
          return Result.Error(
            ErrorOnNode(node: \(node_to_check), withError: "\(raw: error_message)"))
        }
        """)
    ]
  }
}
public struct RequireNodesType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex

    let node_to_check = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    guard let expected_types = node.arguments[arg_index].expression.as(ArrayExprSyntax.self) else {
      throw MacroError(withMessage: "Node(s) to check must be in an array")
    }

    arg_index = arguments.index(after: arg_index)
    guard
      let expected_type_nice_names = node.arguments[arg_index].expression.as(ArrayExprSyntax.self)
    else {
      throw MacroError(withMessage: "Node nice names must be in an array")
    }

    let error_message =
      "Did not find one of the expected types: "
      + expected_type_nice_names.elements.map { l in
        remove_embedded_quotes("\(l.expression)")
      }.joined(separator: ",")

    let ifs = expected_types.elements.map { l in
      "\(node_to_check).nodeType != \(l.expression)"
    }.joined(separator: " && ")

    return [
      CodeBlockItemSyntax(
        """
        if \(raw: ifs) {
          return Result.Error(
            ErrorOnNode(node: \(node_to_check), withError: "\(raw: error_message)"))
        }
        """)
    ]
  }
}
public struct SkipUnlessNodeType: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices
    var arg_index = arguments.startIndex
    let node_to_check = node.arguments[arg_index].expression
    arg_index = arguments.index(after: arg_index)
    let expected_type = node.arguments[arg_index].expression

    return [
      CodeBlockItemSyntax(
        """
        if \(node_to_check).nodeType != \(expected_type) {
          return Result.Ok(.none)
        }
        """)
    ]
  }
}

public struct MustOr: CodeItemMacro {
  public static func expansion(
    of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext
  ) throws -> [CodeBlockItemSyntax] {
    let arguments = node.arguments.indices

    var arg_index = arguments.startIndex
    let result = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    let thing = node.arguments[arg_index].expression

    arg_index = arguments.index(after: arg_index)
    let or = node.arguments[arg_index].expression

    return [
      CodeBlockItemSyntax(
        """
        if let __thing = \(thing) {
          \(result) = __thing 
        } else {
          return \(or)
        }
        """)
    ]
  }
}

@main
struct P4Macros: CompilerPlugin {
  var providingMacros: [Macro.Type] = [
    RequireResult.self, RequireErrorResult.self, UseOkResult.self, UseErrorResult.self,
    RequireNodeType.self, SkipUnlessNodeType.self, RequireNodesType.self, MustOr.self,
  ]
}