Add NodeType-Checking Macros

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

Sources/Macros/Macros.swift

@@ -17,7 +17,23 @@
 
 import SwiftCompilerPlugin
 import SwiftSyntax
-import SwiftSyntaxMacros
+@_spi(ExperimentalLanguageFeature) import SwiftSyntaxMacros
+
+public func remove_embedded_quotes(_ from: String) -> String {
+  return from.replacing("\"", with: [])
+}
+
+struct MacroError: Error, CustomStringConvertible {
+    var message: String
+    var description: String {
+      get {
+        return message 
+      }
+    }
+    public init(withMessage _message: String) {
+      message = _message
+    }
+}
 
 public struct UseOkResult: ExpressionMacro {
   public static func expansion(
@@ -25,7 +41,7 @@ public struct UseOkResult: ExpressionMacro {
     in context: some MacroExpansionContext
   ) throws -> ExprSyntax {
 
-    guard let argument = node.argumentList.first?.expression else {
+    guard let argument = node.arguments.first?.expression else {
       throw Require.Error.SyntaxError
     }
 
@@ -48,7 +64,7 @@ public struct UseErrorResult: ExpressionMacro {
     in context: some MacroExpansionContext
   ) throws -> ExprSyntax {
 
-    guard let argument = node.argumentList.first?.expression else {
+    guard let argument = node.arguments.first?.expression else {
       throw Require.Error.SyntaxError
     }
 
@@ -78,7 +94,7 @@ public struct RequireResult: ExpressionMacro {
     in context: some MacroExpansionContext
   ) throws -> ExprSyntax {
 
-    guard let argument = node.argumentList.first?.expression else {
+    guard let argument = node.arguments.first?.expression else {
       throw Require.Error.SyntaxError
     }
 
@@ -101,9 +117,9 @@ public struct RequireErrorResult: ExpressionMacro {
     in context: some MacroExpansionContext
   ) throws -> ExprSyntax {
 
-    let arguments = node.argumentList.indices
-    let expected_error = node.argumentList[arguments.startIndex].expression
-    let error_producer = node.argumentList[arguments.index(after: arguments.startIndex)].expression
+    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(
       """
@@ -121,9 +137,84 @@ public struct RequireErrorResult: ExpressionMacro {
   }
 }
 
+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)
+      }
+      """)]
+  }
+}
+
 @main
 struct P4Macros: CompilerPlugin {
   var providingMacros: [Macro.Type] = [
-    RequireResult.self, RequireErrorResult.self, UseOkResult.self, UseErrorResult.self,
+    RequireResult.self, RequireErrorResult.self, UseOkResult.self, UseErrorResult.self, RequireNodeType.self, SkipUnlessNodeType.self, RequireNodesType.self
   ]
 }