Working with columns as DuplicateDictionary

Author: wickwirew

Sources/Compiler/Columns.swift

@@ -7,12 +7,14 @@
 
 import OrderedCollections
 
-public typealias Columns = OrderedDictionary<Substring, Type>
+public typealias Columns = DuplicateDictionary<Substring, Type>
 
 extension Columns {
     /// Initializes the columns with their default names that SQLite gives to them.
     init(withDefaultNames types: [Type]) {
         self = types.enumerated()
-            .reduce(into: [:]) { c, v in c["column\(v.offset + 1)"] = v.element }
+            .reduce(into: [:]) { c, v in
+                c.append(v.element, for: "column\(v.offset + 1)")
+            }
     }
 }

Sources/Compiler/Environment.swift

@@ -174,9 +174,13 @@ extension Environment {
     /// Imports the table into the environment.
     /// If `isOptional` is true all columns types will be
     /// forced to their optional value.
+    /// If `qualifiedAccessOnly` is true the columns will only be
+    /// available via qualified access with at least the table
+    /// name specified
     mutating func `import`(
         table: Table,
-        isOptional: Bool
+        isOptional: Bool,
+        qualifiedAccessOnly: Bool = false
     ) {
         let importedTable = ImportedTable(
             table: isOptional ? table.mapTypes { $0.coerceToOptional() } : table,
@@ -189,6 +193,9 @@ extension Environment {
         // in `MATCH` and `IN` statements.
         insert(detached: table.name.name, type: table.type, isOptional: false)
 
+        // Don't insert columns into detached if they required qualified access
+        guard !qualifiedAccessOnly else { return }
+        
         for (column, type) in table.columns {
             insert(detached: column, type: type, isOptional: isOptional)
         }
@@ -267,11 +274,13 @@ extension Environment {
         in importedTable: ImportedTable,
         forceAmbiguous: Bool = false
     ) -> LookupResult<Type> {
-        if let column = importedTable.table.columns[column] {
-            return LookupResult(column, isAmbiguous: forceAmbiguous)
+        let entries = importedTable.table.columns[column]
+        
+        if let column = entries.first {
+            return LookupResult(column, isAmbiguous: forceAmbiguous || entries.count > 1)
         }
 
-        guard let column = importedTable.additionalColumns?[column] else {
+        guard let column = importedTable.additionalColumns?[column].first else {
             // If the table does not have it the parent won't either
             return .columnDoesNotExist(column)
         }

Sources/Compiler/Sema/Statement.swift

@@ -153,7 +153,7 @@ public struct ResultColumns: Sendable {
 
         return chunks.reduce(into: [:]) { result, chunk in
             for column in chunk.columns {
-                result[column.key] = column.value
+                result.append(column.value, for: column.key)
             }
         }
     }

Sources/Compiler/Sema/StmtTypeChecker.swift

@@ -277,7 +277,7 @@ extension StmtTypeChecker: StmtSyntaxVisitor {
             let types = resultColumns.values
             let minCount = min(stmt.columnNames.count, types.count)
             columns = (0..<minCount).reduce(into: [:]) { columns, index in
-                columns[stmt.columnNames[index].value] = types[index]
+                columns.append(types[index], for: stmt.columnNames[index].value)
             }
         }
 
@@ -345,17 +345,17 @@ extension StmtTypeChecker: StmtSyntaxVisitor {
 
         switch stmt.action {
         case .delete:
-            importTable(table, as: "old")
+            importTable(table, as: "old", qualifiedAccessOnly: true)
         case .insert:
-            importTable(table, as: "new")
+            importTable(table, as: "new", qualifiedAccessOnly: true)
         case let .update(columns):
-            importTable(table, as: "new")
-            importTable(table, as: "old")
+            importTable(table, as: "new", qualifiedAccessOnly: true)
+            importTable(table, as: "old", qualifiedAccessOnly: true)
 
             // Make sure all columns in the update statement actually exist
             if let columns {
                 for column in columns {
-                    if table.columns[column.value] == nil {
+                    if !table.columns.contains(key: column.value) {
                         diagnostics.add(.columnDoesNotExist(column))
                     }
                 }
@@ -502,7 +502,7 @@ extension StmtTypeChecker {
         if let columns = insert.columns {
             var columnTypes: [Type] = []
             for column in columns {
-                guard let def = table.columns[column.value] else {
+                guard let def = table.columns[column.value].first else {
                     diagnostics.add(.columnDoesNotExist(column))
                     columnTypes.append(.error)
                     continue
@@ -552,7 +552,7 @@ extension StmtTypeChecker {
             case let .single(column):
                 nameInferrer.suggest(name: column.value, for: valueName)
 
-                guard let column = table.columns[column.value] else {
+                guard let column = table.columns[column.value].first else {
                     diagnostics.add(.columnDoesNotExist(column))
                     return .empty
                 }
@@ -616,7 +616,7 @@ extension StmtTypeChecker {
         var columns: [Type] = []
 
         for name in names {
-            if let column = table.columns[name.value] {
+            if let column = table.columns[name.value].first {
                 columns.append(column)
             } else {
                 diagnostics.add(.columnDoesNotExist(name))
@@ -640,10 +640,9 @@ extension StmtTypeChecker {
 
                 let name = alias?.identifier.value ?? names.proposedName ?? "column\(offset + 1)"
 
-                resultColumns[name] = type
+                resultColumns.append(type, for: name)
             case .all:
-                // TODO: See TODO on `Columns` typealias
-                resultColumns.merge(sourceTable.columns, uniquingKeysWith: { $1 })
+                resultColumns.append(contentsOf: sourceTable.columns)
             }
         }
 
@@ -678,7 +677,7 @@ extension StmtTypeChecker {
             }
 
             columns = (0 ..< min(columnTypes.count, cte.columns.count))
-                .reduce(into: [:]) { $0[cte.columns[$1].value] = columnTypes[$1] }
+                .reduce(into: [:]) { $0.append(columnTypes[$1], for: cte.columns[$1].value) }
         }
 
         return Table(
@@ -796,7 +795,7 @@ extension StmtTypeChecker {
     /// selected all columns from foo and bar so the table structs can be embded
     /// within the output type.
     private mutating func typeCheck(resultColumns: [ResultColumnSyntax]) -> ResultColumns {
-        var columns: OrderedDictionary<Substring, Type> = [:]
+        var columns: Columns = [:]
         var table: Substring?
         // Each chunk is either a list of specifically listed columns
         // or a select all of a table.
@@ -819,7 +818,7 @@ extension StmtTypeChecker {
                 let (type, names) = typeCheck(expr)
                 let name = alias?.identifier.value ?? names.proposedName ?? "column\(offset + 1)"
 
-                columns[name] = type
+                columns.append(type, for: name)
                 nameInferrer.suggest(name: name, for: names)
 
                 // We selected a single column, so clear out the table
@@ -960,17 +959,28 @@ extension StmtTypeChecker {
     private mutating func importTable(
         _ table: Table,
         as alias: Substring? = nil,
-        isOptional: Bool = false
+        isOptional: Bool = false,
+        qualifiedAccessOnly: Bool = false
     ) {
         // Insert real name not alias. These are used later for observation tracking
         // so an alias is no good since it will always be the actual table name.
         usedTableNames.insert(table.name.name)
 
         // Table is always accessible by it's name even if aliased
-        env.import(table: table, isOptional: isOptional)
+        env.import(
+            table: table,
+            isOptional: isOptional,
+            qualifiedAccessOnly: qualifiedAccessOnly
+        )
 
+        // Aliases can only ever be used qualifed since the original
+        // columns are already inserted
         if let alias {
-            env.import(table: table.aliased(to: alias), isOptional: isOptional)
+            env.import(
+                table: table.aliased(to: alias),
+                isOptional: isOptional,
+                qualifiedAccessOnly: true
+            )
         }
     }
 
@@ -994,11 +1004,12 @@ extension StmtTypeChecker {
         case let .columns(columnsDefs, constraints, options):
             var columns: Columns = [:]
             for (name, def) in columnsDefs {
-                columns[name.value] = typeFor(
+                let type = typeFor(
                     column: def,
                     tableColumns: columns,
                     tableName: createTable.name.value
                 )
+                columns.append(type, for: name.value)
             }
 
             validateTableConstraints(
@@ -1053,17 +1064,16 @@ extension StmtTypeChecker {
             tableName = QualifiedName(name: newTableName.value, schema: tableName.schema)
             table.name = tableName
         case let .renameColumn(oldName, newName):
-            table.columns = table.columns.reduce(into: [:]) { newColumns, column in
-                newColumns[column.key == oldName.value ? newName.value : column.key] = column.value
-            }
+            table.columns.rename(oldName.value, to: newName.value)
         case let .addColumn(column):
-            table.columns[column.name.value] = typeFor(
+            let newType = typeFor(
                 column: column,
                 tableColumns: table.columns,
                 tableName: table.name.name
             )
+            table.columns.append(newType, for: column.name.value)
         case let .dropColumn(column):
-            table.columns[column.value] = nil
+            table.columns = Columns(table.columns.filter { $0.key != column.value })
         }
 
         schema[tableName] = table
@@ -1127,13 +1137,13 @@ extension StmtTypeChecker {
                         // declared for so if its this table and this column then ignore it.
                         guard column.name.value != foreignColumn.value else { continue }
 
-                        if tableColumns[foreignColumn.value] == nil {
+                        if !tableColumns.contains(key: foreignColumn.value) {
                             diagnostics.add(.columnDoesNotExist(foreignColumn))
                         }
                     }
                 } else if let table = schema[QualifiedName(name: fk.foreignTable.value, schema: .main)] {
                     for foreignColumn in fk.foreignColumns {
-                        if table.columns[foreignColumn.value] == nil {
+                        if !table.columns.contains(key: foreignColumn.value) {
                             diagnostics.add(.columnDoesNotExist(foreignColumn))
                         }
                     }
@@ -1179,10 +1189,10 @@ extension StmtTypeChecker {
         tableColumns: Columns
     ) {
         for foreignColumn in fk.foreignColumns {
-            // Column constraints can oddly reference themselves
+            // Column constraints can reference themselves
             guard column.name.value != foreignColumn.value else { continue }
 
-            if tableColumns[foreignColumn.value] == nil {
+            if !tableColumns.contains(key: foreignColumn.value) {
                 diagnostics.add(.columnDoesNotExist(foreignColumn))
             }
         }
@@ -1237,7 +1247,7 @@ extension StmtTypeChecker {
                 for column in constraint.0 {
                     guard let name = column.columnName else { continue }
 
-                    if columns[name.value] == nil {
+                    if !columns.contains(key: name.value) {
                         diagnostics.add(.columnDoesNotExist(name))
                     } else {
                         columnNames.append(name.value)
@@ -1263,7 +1273,7 @@ extension StmtTypeChecker {
             case .foreignKey(let fkColumns, let fkClause):
                 // Make sure listed columns exist
                 for column in fkColumns {
-                    guard columns[column.value] == nil else { continue }
+                    guard !columns.contains(key: column.value) else { continue }
                     diagnostics.add(.columnDoesNotExist(column))
                 }
 
@@ -1277,7 +1287,7 @@ extension StmtTypeChecker {
 
                 // Make sure referenced columns exist
                 for column in fkClause.foreignColumns {
-                    guard foreignTable.columns[column.value] == nil else { continue }
+                    guard !foreignTable.columns.contains(key: column.value) else { continue }
                     diagnostics.add(.columnDoesNotExist(column))
                 }
             case .primaryKey, .unique:
@@ -1298,9 +1308,11 @@ extension StmtTypeChecker {
                     continue
                 }
 
-                columns[name.value] = notNull != nil
+                let type: Type = notNull != nil
                     ? .nominal(typeName)
                     : .optional(.nominal(typeName))
+                
+                columns.append(type, for: name.value)
             case .fts5Option:
                 break // Nothing to do, maybe validate these in the future
             case .unknown:

Sources/Compiler/Utils/DuplicateDictionary.swift

@@ -29,8 +29,8 @@ public struct DuplicateDictionary<Key: Hashable, Value> {
 
     /// Tuples cannot be equatable...
     public struct _Element {
-        @usableFromInline let key: Key
-        @usableFromInline let value: Value
+        @usableFromInline var key: Key
+        @usableFromInline var value: Value
 
         @usableFromInline
         @inline(__always)
@@ -119,6 +119,7 @@ public struct DuplicateDictionary<Key: Hashable, Value> {
             var currentIndex = 0
 
             public mutating func next() -> Element? {
+                defer { currentIndex += 1 }
                 switch positions {
                 case .empty:
                     return nil
@@ -145,6 +146,16 @@ public struct DuplicateDictionary<Key: Hashable, Value> {
         self.positions = positions
     }
 
+    public init<S: Sequence>(
+        _ sequence: S
+    ) where S.Element == Element {
+        self = .init()
+        
+        for (k, v) in sequence {
+            append(v, for: k)
+        }
+    }
+    
     /// Appends the value for the given key
     @inline(__always)
     public mutating func append(_ value: Value, for key: Key) {
@@ -182,6 +193,28 @@ public struct DuplicateDictionary<Key: Hashable, Value> {
         return p != .empty
     }
 
+    /// Renames all values under the `key` to the `newKey`
+    public mutating func rename(_ key: Key, to newKey: Key) {
+        guard let positions = positions[key] else { return }
+        
+        for index in positions {
+            _values[index].key = newKey
+        }
+        
+        // Insert positions for new key and remove old one.
+        self.positions[newKey] = positions
+        self.positions[key] = nil
+    }
+    
+    /// Updates all values for the given `key` to be the `newValue`
+    public mutating func updateAll(_ key: Key, to newValue: Value) {
+        guard let positions = positions[key] else { return }
+        
+        for index in positions {
+            _values[index].value = newValue
+        }
+    }
+    
     /// Map over the values and transform them into a new dictionary
     /// with the same keys but with the transformed values.
     public func mapValues<T>(
@@ -334,3 +367,9 @@ extension DuplicateDictionary: Sendable where Key: Sendable, Value: Sendable {}
 extension DuplicateDictionary._Element: Sendable where Key: Sendable, Value: Sendable {}
 extension DuplicateDictionary._Element: Equatable where Key: Equatable, Value: Equatable {}
 extension DuplicateDictionary: Equatable where Key: Equatable, Value: Equatable {}
+
+extension DuplicateDictionary: CustomReflectable {
+    public var customMirror: Mirror {
+        Mirror(self, unlabeledChildren: self.map(\.self), displayStyle: .dictionary)
+    }
+}

Tests/CompilerTests/CompilerTests.swift

@@ -31,7 +31,7 @@ class CompilerTests: XCTestCase {
     }
 
     func testCreateTable() throws {
-        try checkSchema(compile: "CompileCreateTable")
+        try checkSchema(compile: "CompileCreateTable", dump: true)
     }
 
     func testDropTable() throws {