Allow for custom input and output names

Author: wickwirew

Sources/Compiler/Compiler.swift

@@ -113,7 +113,7 @@ fileprivate struct CompilerWithSource {
         self.schema = typeChecker.schema
 
         let statement = Statement(
-            name: nil,
+            definition: nil,
             parameters: uniqueParameters,
             resultColumns: type,
             outputCardinality: cardinality,
@@ -190,7 +190,14 @@ extension CompilerWithSource: StmtSyntaxVisitor {
 
     mutating func visit(_ stmt: QueryDefinitionStmtSyntax) -> (Statement, Diagnostics)? {
         guard let (innerStmt, diagnostics) = stmt.statement.accept(visitor: &self) else { return nil }
-        return (innerStmt.with(name: stmt.name.value), diagnostics)
+        
+        let definition = Definition(
+            name: stmt.name.value,
+            input: stmt.input?.value,
+            output: stmt.output?.value
+        )
+        
+        return (innerStmt.with(definition: definition), diagnostics)
     }
 
     mutating func visit(_ stmt: PragmaStmt) -> (Statement, Diagnostics)? {

Sources/Compiler/Gen/Language.swift

@@ -64,12 +64,12 @@ extension Language {
         for statement: Statement,
         tables: OrderedDictionary<Substring, GeneratedModel>
     ) -> GeneratedQuery {
-        guard let name = statement.name else {
+        guard let definition = statement.definition else {
             fatalError("Upstream error should have caught this")
         }
 
-        let input = inputTypeIfNeeded(statement: statement, name: name)
-        let output = outputTypeIfNeeded(statement: statement, name: name, tables: tables)
+        let input = inputTypeIfNeeded(statement: statement, definition: definition)
+        let output = outputTypeIfNeeded(statement: statement, definition: definition, tables: tables)
 
         let type = queryType(
             for: statement.noOutput ? nil : statement.outputCardinality,
@@ -91,7 +91,7 @@ extension Language {
         }.joined()
 
         return GeneratedQuery(
-            name: "\(name)Query",
+            name: "\(definition.name)Query",
             type: type,
             input: input,
             output: output,
@@ -119,7 +119,7 @@ extension Language {
 
     private static func inputTypeIfNeeded(
         statement: Statement,
-        name: Substring
+        definition: Definition
     ) -> BuiltinOrGenerated? {
         guard let firstParameter = statement.parameters.first else { return nil }
 
@@ -130,7 +130,7 @@ extension Language {
             )
         }
 
-        let inputTypeName = "\(name.capitalizedFirst)Input"
+        let inputTypeName = definition.input?.description ?? "\(definition.name.capitalizedFirst)Input"
 
         let model = GeneratedModel(
             name: inputTypeName,
@@ -149,7 +149,7 @@ extension Language {
 
     private static func outputTypeIfNeeded(
         statement: Statement,
-        name: Substring,
+        definition: Definition,
         tables: OrderedDictionary<Substring, GeneratedModel>
     ) -> BuiltinOrGenerated? {
         guard let firstResultColumns = statement.resultColumns.chunks.first else { return nil }
@@ -171,7 +171,7 @@ extension Language {
             return .builtin(builtinType(for: firstColumn), isArray: firstColumn.isRow)
         }
 
-        let outputTypeName = "\(name.capitalizedFirst)Output"
+        let outputTypeName = definition.output?.description ?? "\(definition.name.capitalizedFirst)Output"
 
         let model = GeneratedModel(
             name: outputTypeName,

Sources/Compiler/Parse/Parsers.swift

@@ -2070,11 +2070,28 @@ enum Parsers {
         let define = state.take(.define)
         state.skip(.query)
         let name = identifier(state: &state)
+        
+        let params: [Substring: IdentifierSyntax]? = take(if: .openParen, state: &state) { state in
+            parens(state: &state) { state in
+                delimited(by: .comma, state: &state, reduceInto: [:]) { params, param in
+                    params[param.name.value] = param.value
+                } element: { state in
+                    let name = identifier(state: &state)
+                    state.skip(.colon)
+                    let value = identifier(state: &state)
+                    return (name: name, value: value)
+                }
+            }
+        }
+        
         state.skip(.as)
         let stmt = try stmt(state: &state)
+        
         return QueryDefinitionStmtSyntax(
             id: state.nextId(),
             name: name,
+            input: params?["input"],
+            output: params?["output"],
             statement: stmt,
             location: define.location.spanning(state.current.location)
         )
@@ -2143,6 +2160,22 @@ enum Parsers {
         return elements
     }
 
+    static func delimited<Element, Output>(
+        by kind: Token.Kind,
+        state: inout ParserState,
+        reduceInto elements: Output,
+        reduce: (inout Output, Element) -> Void,
+        element: (inout ParserState) throws -> Element
+    ) rethrows -> Output {
+        var elements = elements
+
+        repeat {
+            try reduce(&elements, element(&state))
+        } while state.take(if: kind)
+        
+        return elements
+    }
+    
     static func parens<Value>(
         state: inout ParserState,
         value: (inout ParserState) throws -> Value

Sources/Compiler/Sema/Statement.swift

@@ -6,7 +6,8 @@
 //
 
 public struct Statement {
-    public let name: Substring?
+    /// The information in the `DEFINE` statement if it exists
+    public let definition: Definition?
     /// Any bind parameters for the statement
     public let parameters: [Parameter<String>]
     /// The return type if any.
@@ -29,10 +30,15 @@ public struct Statement {
         return resultColumns.isEmpty
     }
 
-    /// Replaces the name with the given input
-    public func with(name: Substring?) -> Statement {
+    /// The name if one was defined
+    public var name: Substring? {
+        return definition?.name
+    }
+    
+    /// Replaces the definition with the given input
+    public func with(definition: Definition?) -> Statement {
         return Statement(
-            name: name,
+            definition: definition,
             parameters: parameters,
             resultColumns: resultColumns,
             outputCardinality: outputCardinality,
@@ -65,6 +71,13 @@ public struct Parameter<Name> {
     }
 }
 
+/// The values from a `DEFINE QUERY` statement
+public struct Definition {
+    public let name: Substring
+    public let input: Substring?
+    public let output: Substring?
+}
+
 /// The output of a statement
 public struct ResultColumns: Sendable {
     public let chunks: [Chunk]

Sources/Compiler/Syntax/StmtSyntax.swift

@@ -115,6 +115,8 @@ struct InsertStmtSyntax: StmtSyntax, Syntax {
 struct QueryDefinitionStmtSyntax: StmtSyntax {
     let id: SyntaxId
     let name: IdentifierSyntax
+    let input: IdentifierSyntax?
+    let output: IdentifierSyntax?
     let statement: any StmtSyntax
     let location: SourceLocation
 

Tests/CompilerTests/Parser/ParseDefinition.sql

@@ -12,3 +12,20 @@ SELECT * FROM user WHERE id = ?;
 -- CHECK:   ...
 DEFINE QUERY insertUser AS
 INSERT INTO user (id, name) VALUES (1, 'Joe');
+
+-- CHECK: QUERY_DEFINITION_STMT_SYNTAX
+-- CHECK:   NAME fetchUser
+-- CHECK:   OUTPUT FetchedUser
+-- CHECK:   STATEMENT
+-- CHECK:   ...
+DEFINE QUERY fetchUser(output: FetchedUser) AS
+SELECT * FROM user WHERE id = ?;
+
+-- CHECK: QUERY_DEFINITION_STMT_SYNTAX
+-- CHECK:   NAME fetchUser
+-- CHECK:   INPUT TheBestInput
+-- CHECK:   OUTPUT FetchedUser
+-- CHECK:   STATEMENT
+-- CHECK:   ...
+DEFINE QUERY fetchUser(input: TheBestInput, output: FetchedUser) AS
+SELECT * FROM user WHERE id = ?;