diff --git a/runtime/Swift/Sources/Antlr4/atn/ATNConfigSet.swift b/runtime/Swift/Sources/Antlr4/atn/ATNConfigSet.swift
index 66e650237..15254b350 100644
--- a/runtime/Swift/Sources/Antlr4/atn/ATNConfigSet.swift
+++ b/runtime/Swift/Sources/Antlr4/atn/ATNConfigSet.swift
@@ -321,17 +321,17 @@ public class ATNConfigSet: Hashable, CustomStringConvertible {
         return Array(configToAlts.values)
     }
 
-    public final func getStateToAltMap() -> [ATNState: BitSet] {
+    public func getStateToAltMap() -> [Int: BitSet] {
         let length = configs.count
-        var m = [ATNState: BitSet]()
+        var m = [Int: BitSet]()
 
         for i in 0..<length {
             var alts: BitSet
-            if let mAlts =  m[configs[i].state] {
+            if let mAlts =  m[configs[i].state.stateNumber] {
                 alts = mAlts
             } else {
                 alts = BitSet()
-                m[configs[i].state] = alts
+                m[configs[i].state.stateNumber] = alts
             }
 
             try! alts.set(configs[i].alt)
diff --git a/runtime/Swift/Sources/Antlr4/atn/PredictionMode.swift b/runtime/Swift/Sources/Antlr4/atn/PredictionMode.swift
index 31db58404..1d4e9ed05 100644
--- a/runtime/Swift/Sources/Antlr4/atn/PredictionMode.swift
+++ b/runtime/Swift/Sources/Antlr4/atn/PredictionMode.swift
@@ -485,20 +485,8 @@ public enum PredictionMode {
         return configs.getConflictingAltSubsets()
     }
 
-    /// 
-    /// Get a map from state to alt subset from a configuration set. For each
-    /// configuration `c` in `configs`:
-    /// 
-    /// 
-    /// map[c._org.antlr.v4.runtime.atn.ATNConfig#state state_] U= c._org.antlr.v4.runtime.atn.ATNConfig#alt alt_
-    /// 
-    /// 
-    public static func getStateToAltMap(_ configs: ATNConfigSet) -> [ATNState: BitSet] {
-        return configs.getStateToAltMap()
-    }
-
     public static func hasStateAssociatedWithOneAlt(_ configs: ATNConfigSet) -> Bool {
-        let x = getStateToAltMap(configs)
+        let x = configs.getStateToAltMap()
         for alts in x.values {
             if alts.cardinality() == 1 {
                 return true