Finite state networks: replace some lookup tables with maps. Rewrite closure for speed.

2 files changed, 17 insertions, 9 deletions

diff --git a/src/GF/Speech/CFGToFiniteState.hs b/src/GF/Speech/CFGToFiniteState.hs
index 25790786a..5f8e3a093 100644
--- a/src/GF/Speech/CFGToFiniteState.hs
+++ b/src/GF/Speech/CFGToFiniteState.hs
@@ -130,8 +130,10 @@ make_fa c@(g,ns) q0 alpha q1 fa =
                                     fa''' = foldl (\f (CFRule c (Cat d:xs) _) -> make_fa_ (getState d) xs (getState c) f) fa'' rs
                                 in newTransition (getState a) q1 Nothing fa'''
                           where
-                          (fa',ss) = addStatesForCats ni fa
-                          getState x = lookup' x ss
+                          (fa',stateMap) = addStatesForCats ni fa
+                          getState x = Map.findWithDefault 
+                                       (error $ "CFGToFiniteState: No state for " ++ x) 
+                                       x stateMap
                         -- a is not recursive
                         Nothing -> let rs = catRules g a
                                     in foldl (\fa -> \ (CFRule _ b _) -> make_fa_ q0 b q1 fa) fa rs
@@ -140,9 +142,10 @@ make_fa c@(g,ns) q0 alpha q1 fa =
   where
   make_fa_ = make_fa c
 
-addStatesForCats :: [Cat_] -> NFA Token -> (NFA Token, [(Cat_,State)])
-addStatesForCats cs fa = (fa', zip cs (map fst ns))
+addStatesForCats :: [Cat_] -> NFA Token -> (NFA Token, Map Cat_ State)
+addStatesForCats cs fa = (fa', m)
   where (fa', ns) = newStates (replicate (length cs) ()) fa
+        m = Map.fromList (zip cs (map fst ns))
 
 ruleIsNonRecursive :: Set Cat_ -> CFRule_ -> Bool
 ruleIsNonRecursive cs = noCatsInSet cs . ruleRhs
diff --git a/src/GF/Speech/FiniteState.hs b/src/GF/Speech/FiniteState.hs
index e3975d498..e48064945 100644
--- a/src/GF/Speech/FiniteState.hs
+++ b/src/GF/Speech/FiniteState.hs
@@ -152,16 +152,21 @@ determinize (FA g s f) = let (ns,es) = h [start] [] []
 numberStates :: (Ord x,Enum y) => FA x a b -> FA y a b
 numberStates (FA g s fs) = FA (renameNodes newName rest g) s' fs'
     where (ns,rest) = splitAt (length (nodes g)) $ [toEnum 0 .. ]
-          newNodes = zip (map fst (nodes g)) ns
-	  newName n = lookup' n newNodes
+          newNodes = Map.fromList (zip (map fst (nodes g)) ns)
+	  newName n = Map.findWithDefault (error "FiniteState.newName") n newNodes
           s' = newName s
           fs' = map newName fs
 
 -- | Get all the nodes reachable from a list of nodes by only empty edges.
 closure :: Ord n => Outgoing n a (Maybe b) -> Set n -> Set n
-closure out = fix closure_
-    where closure_ r = inserts [y | x <- Set.toList r, (_,y,Nothing) <- getOutgoing out x] r
-          inserts xs s = foldl (flip Set.insert) s xs
+closure out x = closure_ x x
+  where closure_ acc check | Set.null check = acc
+                           | otherwise = closure_ acc' check'
+            where
+            reach = Set.fromList [y | x <- Set.toList check, 
+                                      (_,y,Nothing) <- getOutgoing out x]
+            acc' = acc `Set.union` reach
+            check' = reach Set.\\ acc
 
 -- | Get a map of labels to sets of all nodes reachable 
 --   from a the set of nodes by one edge with the given