feature: Adds visitor editing support

Author: NeedleInAJayStack

Sources/GraphQL/Language/AST.swift

@@ -174,6 +174,28 @@ public enum NodeResult {
             return .node(array[key])
         }
     }
+
+    func set(value: NodeResult, key: IndexPathElement) -> Self? {
+        switch self {
+        case let .node(node):
+            guard let key = key.keyValue else {
+                return nil
+            }
+            node.set(value: value, key: key)
+            return .node(node)
+        case var .array(array):
+            switch value {
+            case let .node(value):
+                guard let key = key.indexValue else {
+                    return nil
+                }
+                array[key] = value
+                return .array(array)
+            case let .array(value):
+                return .array(value)
+            }
+        }
+    }
 }
 
 /**
@@ -183,15 +205,17 @@ public protocol Node {
     var kind: Kind { get }
     var loc: Location? { get }
     func get(key: String) -> NodeResult?
-    func set(value: Node?, key: String)
+    func set(value: NodeResult?, key: String)
 }
 
 public extension Node {
     func get(key _: String) -> NodeResult? {
         return nil
     }
 
-    func set(value _: Node?, key _: String) {}
+    func set(value: NodeResult?, key: String) {
+        print("TODO: Should be implemented on each type!")
+    }
 }
 
 extension Name: Node {}
@@ -252,7 +276,7 @@ extension Name: Equatable {
 public final class Document {
     public let kind: Kind = .document
     public let loc: Location?
-    public let definitions: [Definition]
+    public var definitions: [Definition]
 
     init(loc: Location? = nil, definitions: [Definition]) {
         self.loc = loc
@@ -270,6 +294,24 @@ public final class Document {
             return nil
         }
     }
+    
+    public func set(value: NodeResult?, key: String) {
+        guard let value = value else {
+            return
+        }
+        switch key {
+        case "definitions":
+            guard
+                case let .array(values) = value,
+                let definitions = values as? [Definition]
+            else {
+                return
+            }
+            self.definitions = definitions
+        default:
+            return
+        }
+    }
 }
 
 extension Document: Equatable {
@@ -323,10 +365,10 @@ public final class OperationDefinition {
     public let kind: Kind = .operationDefinition
     public let loc: Location?
     public let operation: OperationType
-    public let name: Name?
-    public let variableDefinitions: [VariableDefinition]
-    public let directives: [Directive]
-    public let selectionSet: SelectionSet
+    public var name: Name?
+    public var variableDefinitions: [VariableDefinition]
+    public var directives: [Directive]
+    public var selectionSet: SelectionSet
 
     init(
         loc: Location? = nil,
@@ -364,6 +406,48 @@ public final class OperationDefinition {
             return nil
         }
     }
+    
+    public func set(value: NodeResult?, key: String) {
+        guard let value = value else {
+            return
+        }
+        switch key {
+        case "name":
+            guard
+                case let .node(node) = value,
+                let name = node as? Name
+            else {
+                return
+            }
+            self.name = name
+        case "variableDefinitions":
+            guard
+                case let .array(values) = value,
+                let variableDefinitions = values as? [VariableDefinition]
+            else {
+                return
+            }
+            self.variableDefinitions = variableDefinitions
+        case "directives":
+            guard
+                case let .array(values) = value,
+                let directives = values as? [Directive]
+            else {
+                return
+            }
+            self.directives = directives
+        case "selectionSet":
+            guard
+                case let .node(value) = value,
+                let selectionSet = value as? SelectionSet
+            else {
+                return
+            }
+            self.selectionSet = selectionSet
+        default:
+            return
+        }
+    }
 }
 
 extension OperationDefinition: Hashable {

Sources/GraphQL/Language/Visitor.swift

@@ -123,14 +123,22 @@ func visit(root: Node, visitor: Visitor, keyMap: [Kind: [String]] = [:]) -> Node
                         }
                     }
                 } else {
-                    let clone = node
-                    node = clone
-                    for (editKey, editValue) in edits {
-                        if case .node(let node) = node {
-                            node.set(value: editValue, key: editKey.keyValue!)
+                    if case let .node(n) = node {
+                        for (editKey, editValue) in edits {
+                            if let editValue = editValue {
+                                if let key = editKey.keyValue {
+                                    n.set(value: .node(editValue), key: key)
+                                }
+                            }
                         }
+                        node = .node(n)
                     }
                 }
+
+                // Since Swift cannot mutate node in-place, we must pass the changes up to parent.
+                if let key = key, let node = node {
+                    parent = parent?.set(value: node, key: key)
+                }
             }
 
             index = stack!.index
@@ -139,15 +147,8 @@ func visit(root: Node, visitor: Visitor, keyMap: [Kind: [String]] = [:]) -> Node
             inArray = stack!.inArray
             stack = stack!.prev
         } else if let parent = parent {
-            key = inArray ? index : keys[index]
-
-            switch parent {
-            case let .node(parent):
-                node = parent.get(key: key!.keyValue!)
-            case let .array(parent):
-                node = .node(parent[key!.indexValue!])
-            }
-
+            key = keys[index]
+            node = parent.get(key: key!)
             if node == nil {
                 continue
             }
@@ -176,9 +177,7 @@ func visit(root: Node, visitor: Visitor, keyMap: [Kind: [String]] = [:]) -> Node
 
             if case .break = result {
                 break
-            }
-
-            if case .skip = result {
+            } else if case .skip = result {
                 if !isLeaving {
                     _ = path.popLast()
                     continue
@@ -204,12 +203,13 @@ func visit(root: Node, visitor: Visitor, keyMap: [Kind: [String]] = [:]) -> Node
             _ = path.popLast()
         } else {
             stack = Stack(index: index, keys: keys, edits: edits, inArray: inArray, prev: stack)
-            inArray = node!.isArray
             switch node! {
             case let .node(node):
+                inArray = false
                 keys = visitorKeys[node.kind] ?? []
             case let .array(array):
-                keys = array.map { _ in "root" }
+                inArray = true
+                keys = Array(array.indices) // array.map { _ in "root" }
             }
             index = -1
             edits = []
@@ -225,7 +225,12 @@ func visit(root: Node, visitor: Visitor, keyMap: [Kind: [String]] = [:]) -> Node
         return nextEditNode
     }
 
-    return root
+    switch node! {
+    case let .node(root): // This should be equal to root, with any relevant edits
+        return root
+    case let .array(array): // This should never occur
+        return array[0]
+    }
 }
 
 final class Stack {