Add schema functionality to directives

Author: kabiroberai

Sources/GenerateTestFixtures/main.swift

@@ -1,19 +1,25 @@
 import Foundation
 import OrionProcessor
 
-// NOTE: These fixtures double as the glue file for Hooks.x.swift
+// NOTE: These fixtures double as the glue file for the runtime tests
 
-let engine = OrionDiagnosticEngine()
-engine.addConsumer(.printing)
+// we want to scope this so that the objects' deinits are called
+do {
+    let engine = OrionDiagnosticEngine()
+    engine.addConsumer(.printing)
 
-let orionTests = URL(fileURLWithPath: #file).deletingLastPathComponent().appendingPathComponent("../../Tests/OrionTests")
+    let orionTests = URL(fileURLWithPath: #filePath)
+        .deletingLastPathComponent()
+        .appendingPathComponent("../../Tests/OrionTests")
 
-let data = try OrionBatchParser(inputs: [orionTests], diagnosticEngine: engine).parse()
-let options = OrionGenerator.Options(emitSourceLocations: false)
-let generator = OrionGenerator(data: data, diagnosticEngine: engine, options: options)
-let glue = try generator.generate()
+    let parserOptions = OrionParser.Options()
+    let data = try OrionBatchParser(inputs: [orionTests], diagnosticEngine: engine, options: parserOptions).parse()
+    let generatorOptions = OrionGenerator.Options(emitSourceLocations: false)
+    let generator = OrionGenerator(data: data, diagnosticEngine: engine, options: generatorOptions)
+    let glue = try generator.generate()
 
-let dest = orionTests.appendingPathComponent("Generated.xc.swift")
-try glue.write(to: dest, atomically: true, encoding: .utf8)
+    let dest = orionTests.appendingPathComponent("Generated.xc.swift")
+    try glue.write(to: dest, atomically: true, encoding: .utf8)
 
-print("Wrote glue file to \(dest.standardized.path)")
+    print("Wrote glue file to \(dest.standardized.path)")
+}

Sources/OrionProcessor/OrionBatchParser.swift

@@ -33,10 +33,12 @@ public final class OrionBatchParser {
 
     public let inputs: [URL]
     public let diagnosticEngine: OrionDiagnosticEngine
+    public let options: OrionParser.Options
     // Note: inputs can include directories
-    public init(inputs: [URL], diagnosticEngine: OrionDiagnosticEngine = .init()) {
+    public init(inputs: [URL], diagnosticEngine: OrionDiagnosticEngine = .init(), options: OrionParser.Options = .init()) {
         self.inputs = inputs
         self.diagnosticEngine = diagnosticEngine
+        self.options = options
     }
 
     private func computeFiles() throws -> [URL] {
@@ -71,7 +73,7 @@ public final class OrionBatchParser {
         let files = try computeFiles().sorted { $0.path < $1.path }
         let engine = diagnosticEngine
         let allData = files.concurrentMap { file -> Result<OrionData, Error> in
-            Result { try OrionParser(file: file, diagnosticEngine: engine).parse() }
+            Result { try OrionParser(file: file, diagnosticEngine: engine, options: options).parse() }
         }
         return OrionData(merging: try allData.map { try $0.get() })
     }

Sources/OrionProcessor/OrionDirective.swift

@@ -78,7 +78,7 @@ struct OrionDirectiveDiagnostic: Error, LocalizedError {
 
 final class OrionDirectiveParser {
     static let shared = OrionDirectiveParser()
-    private static let prefix = "orion:"
+    private static let prefix = "orion"
 
     private let exactMap: [String: OrionDirective.Type]
     private let prefixMap: [String: OrionDirective.Type]
@@ -115,10 +115,29 @@ final class OrionDirectiveParser {
         return customList.lazy.compactMap { $0(name) }.first
     }
 
-    func directive(from text: String, at location: SourceLocation) throws -> OrionDirective? {
+    func directive(from text: String, at location: SourceLocation, schema: Set<String> = []) throws -> OrionDirective? {
         guard text.hasPrefix(Self.prefix) else { return nil }
         let dropped = text.dropFirst(Self.prefix.count)
-        let parts = dropped.split(separator: " ")
+        // directives can be of the form `orion:foo` or `orion[mySchema]:foo`.
+        // check which form it is, and if it's the latter then only parse
+        // the directive if the schema is within our enabled schema set.
+        let body: Substring // the `foo` bit
+        switch dropped.first {
+        case ":":
+            body = dropped.dropFirst()
+        case "[":
+            let withoutOpening = dropped.dropFirst() // mySchema]:foo
+            guard let endIdx = withoutOpening.firstIndex(of: "]")
+            else { return nil }
+            let textSchema = withoutOpening[..<endIdx] // mySchema
+            guard schema.contains(String(textSchema)) else { return nil }
+            let afterSchema = withoutOpening[endIdx...] // ]:foo
+            guard afterSchema.hasPrefix("]:") else { return nil }
+            body = afterSchema.dropFirst(2)
+        default:
+            return nil
+        }
+        let parts = body.split(separator: " ")
         guard let name = parts.first.map(String.init) else { return nil }
         let arguments = parts.dropFirst().map(String.init)
         guard let matched = type(matching: name) else {
@@ -147,7 +166,8 @@ enum OrionDirectives {
         ReturnsRetained.self,
         New.self,
         Disable.self,
-        SuprTramp.self
+        SuprTramp.self,
+        IgnoreImport.self,
     ]
 
     struct ReturnsRetained: OrionDirective {
@@ -208,4 +228,16 @@ enum OrionDirectives {
             }
         }
     }
+
+    struct IgnoreImport: OrionDirective {
+        static let matchRule: OrionDirectiveMatchRule = .exact("ignore_import")
+
+        let base: OrionDirectiveBase
+        init(base: OrionDirectiveBase) throws {
+            self.base = base
+            guard base.arguments.isEmpty else {
+                throw OrionDirectiveDiagnostic("ignore_import directive expected zero arguments, got \(base.arguments.count)")
+            }
+        }
+    }
 }

Sources/OrionProcessor/OrionParser.swift

@@ -3,6 +3,14 @@ import SwiftSyntax
 
 public final class OrionParser {
 
+    public struct Options {
+        public let schema: Set<String>
+
+        public init(schema: Set<String> = []) {
+            self.schema = schema
+        }
+    }
+
     private enum Source {
         case file(URL)
         case contents(String)
@@ -25,22 +33,25 @@ public final class OrionParser {
     }
 
     public let engine: DiagnosticEngine
+    public let options: Options
     private let source: Source
 
-    public init(file: URL, diagnosticEngine: OrionDiagnosticEngine = .init()) {
+    public init(file: URL, diagnosticEngine: OrionDiagnosticEngine = .init(), options: Options = .init()) {
         source = .file(file)
         self.engine = diagnosticEngine.createEngine()
+        self.options = options
     }
 
-    public init(contents: String, diagnosticEngine: OrionDiagnosticEngine = .init()) {
+    public init(contents: String, diagnosticEngine: OrionDiagnosticEngine = .init(), options: Options = .init()) {
         source = .contents(contents)
         self.engine = diagnosticEngine.createEngine()
+        self.options = options
     }
 
     public func parse() throws -> OrionData {
         let syntax = try source.parseSyntax(diagnosticEngine: engine)
         let converter = SourceLocationConverter(file: source.filename, tree: syntax)
-        let visitor = OrionVisitor(diagnosticEngine: engine, sourceLocationConverter: converter)
+        let visitor = OrionVisitor(diagnosticEngine: engine, sourceLocationConverter: converter, options: options)
         visitor.walk(syntax)
         guard !visitor.didFail else {
             throw OrionFailure()

Sources/OrionProcessor/OrionVisitor.swift

@@ -146,11 +146,17 @@ class OrionVisitor: SyntaxVisitor {
         "hookWillActivate", "hookDidActivate"
     ]
 
+    let options: OrionParser.Options
     let converter: SourceLocationConverter
     let diagnosticEngine: DiagnosticEngine
-    init(diagnosticEngine: DiagnosticEngine, sourceLocationConverter: SourceLocationConverter) {
+    init(
+        diagnosticEngine: DiagnosticEngine,
+        sourceLocationConverter: SourceLocationConverter,
+        options: OrionParser.Options
+    ) {
         self.diagnosticEngine = diagnosticEngine
         self.converter = sourceLocationConverter
+        self.options = options
     }
 
     private(set) var data = OrionData()
@@ -185,7 +191,7 @@ class OrionVisitor: SyntaxVisitor {
                 return nil
             }
             do {
-                return try OrionDirectiveParser.shared.directive(from: directive, at: location)
+                return try OrionDirectiveParser.shared.directive(from: directive, at: location, schema: options.schema)
             } catch let err as OrionDirectiveDiagnostic {
                 if warnOnFailure {
                     diagnosticEngine.diagnose(err.diagnosticMessage, location: location)
@@ -601,6 +607,12 @@ class OrionVisitor: SyntaxVisitor {
     }
 
     override func visitPost(_ node: ImportDeclSyntax) {
+        let directives = makeDirectives(for: Syntax(node))
+        let ignoreImports = directives.filter { $0 is OrionDirectives.IgnoreImport }
+        guard ignoreImports.isEmpty else {
+            ignoreImports.forEach { $0.setUsed() }
+            return
+        }
         data.imports.append(node.withoutTrivia())
     }
 

Sources/OrionProcessorCLI/main.swift

@@ -53,6 +53,17 @@ struct OrionCommand: ParsableCommand {
         )
     ) var backendModules: [String] = []
 
+    @Option(
+        help: ArgumentHelp(
+            "Add name to schema.",
+            discussion: """
+            The schema list can be used to conditionally enable Orion directives. It is \
+            empty by default. Read the documentation on directives for more information.
+            """,
+            valueName: "name"
+        )
+    ) var schema: [String] = []
+
     @Flag(
         inversion: .prefixedNo,
         help: ArgumentHelp(
@@ -85,6 +96,8 @@ struct OrionCommand: ParsableCommand {
         let engine = OrionDiagnosticEngine()
         engine.addConsumer(.printing)
 
+        let parserOptions = OrionParser.Options(schema: Set(schema))
+
         let data: OrionData
         if inputs == ["-"] {
             var input = ""
@@ -94,23 +107,25 @@ struct OrionCommand: ParsableCommand {
             }
             let parser = OrionParser(
                 contents: input,
-                diagnosticEngine: engine
+                diagnosticEngine: engine,
+                options: parserOptions
             )
             data = try parser.parse()
         } else {
             let parser = OrionBatchParser(
                 inputs: inputs.map(URL.init(fileURLWithPath:)),
-                diagnosticEngine: engine
+                diagnosticEngine: engine,
+                options: parserOptions
             )
             data = try parser.parse()
         }
 
-        let options = OrionGenerator.Options(
+        let generatorOptions = OrionGenerator.Options(
             backend: backend,
             extraBackendModules: Set(backendModules),
             emitSourceLocations: sourceLocations
         )
-        let generator = OrionGenerator(data: data, diagnosticEngine: engine, options: options)
+        let generator = OrionGenerator(data: data, diagnosticEngine: engine, options: generatorOptions)
         let out = try generator.generate()
         if let output = output {
             let outputURL = URL(fileURLWithPath: output)