wta: Fix whitespace and refactor generator

src/random-wta/Generator.hs

@@ -2,7 +2,13 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE LambdaCase #-}
 
-module Generator (genWTA, runGenerator, GeneratorConfig(..), EdgeConfig(..)) where
+module Generator
+  ( genWTA
+  , runGenerator
+  , GeneratorConfig(..)
+  , EdgeConfig(..)
+  )
+where
 
 import           Data.Vector                    ( Vector )
 import qualified Data.Vector                   as V
@@ -12,8 +18,8 @@ import           Data.Coerce
 import           Data.Maybe
 import           Data.Foldable
 import           Control.Arrow                  ( (&&&) )
-import qualified Data.Map.Strict            as M
-import qualified Data.Set                   as S
+import qualified Data.Map.Strict               as M
+import qualified Data.Set                      as S
 import           Data.Coerce
 import           Data.Ratio
 import           System.IO
@@ -93,61 +99,82 @@ genStateTransitions = do
   arities <- asks (numSymbols . spec)
   fold <$> traverse genForArity (V.findIndices (/= 0) arities)
 
-genTransitions :: Generator m (Vector (Vector (Transition m)))
-genTransitions = do
+genTransitionsZeroFreq :: Generator m (Vector (Vector (Transition m)))
+genTransitionsZeroFreq = do
   n <- asks (numStates . spec)
   V.replicateM n genStateTransitions
 
 -- TODO Implement (Random IndexedTransition)
 uniqueTransitions :: Int -> IndexedTransition -> IO [IndexedTransition]
 uniqueTransitions num idxMax@(IndexedTransition.Index max)
-  | fromIntegral num < fromIntegral max * (7%10) = uniqueTransitionsByGeneration num idxMax
-  | otherwise       = uniqueTransitionsByElimination num idxMax
-
-uniqueTransitionsByGeneration :: Int -> IndexedTransition -> IO [IndexedTransition]
-uniqueTransitionsByGeneration num (IndexedTransition.Index max) = helper S.empty num
+  | fromIntegral num < fromIntegral max * (7 % 10) = uniqueTransitionsByGeneration
+    num
+    idxMax
+  | otherwise = uniqueTransitionsByElimination num idxMax
+
+uniqueTransitionsByGeneration
+  :: Int -> IndexedTransition -> IO [IndexedTransition]
+uniqueTransitionsByGeneration num (IndexedTransition.Index max) = helper
+  S.empty
+  num
   where
     helper m 0 = return $ coerce (S.toList m)
     helper m c = do
-      x <- randomRIO (0, max-1)
-      if x `S.member` m then helper m c else helper (S.insert x m) (c-1)
-
-uniqueTransitionsByElimination :: Int -> IndexedTransition -> IO [IndexedTransition]
-uniqueTransitionsByElimination num (IndexedTransition.Index max) = helper whole num
+      x <- randomRIO (0, max - 1)
+      if x `S.member` m then helper m c else helper (S.insert x m) (c - 1)
+
+uniqueTransitionsByElimination
+  :: Int -> IndexedTransition -> IO [IndexedTransition]
+uniqueTransitionsByElimination num (IndexedTransition.Index max) = helper
+  whole
+  num
   where
     helper free 0 = return $ coerce (S.toList (S.difference whole free))
     helper free c = do
       idx <- randomRIO (0, S.size free - 1)
       let x = S.elemAt idx free
-      helper (S.delete x free) (c-1)
+      helper (S.delete x free) (c - 1)
 
-    whole = S.fromList [0..max-1]
+    whole = S.fromList [0 .. max - 1]
 
-genTransitions' :: Int -> Generator m (Vector (Vector (Transition m)))
-genTransitions' numTransitions = do
+genTransitionsNumTrans :: Int -> Generator m (Vector (Vector (Transition m)))
+genTransitionsNumTrans numTransitions = do
   wtaSpec <- asks spec
-  let n = numStates wtaSpec
+  let n    = numStates wtaSpec
       maxT = IndexedTransition.maxIndex wtaSpec
 
-  numTransitions' <- if (fromIntegral numTransitions > maxT) then do
+  numTransitions' <- if (fromIntegral numTransitions > maxT)
+    then do
       let cap = IndexedTransition.fromIndexd maxT
-      lift $ hPutStrLn stderr ("warning: More transitions than possible requested. Capping at " <> show cap)
+      lift $ hPutStrLn
+        stderr
+        (  "warning: More transitions than possible requested. Capping at "
+        <> show cap
+        )
       return (fromIntegral cap)
     else do
       return numTransitions
 
-  transitions <- lift $ map (IndexedTransition.fromIndex wtaSpec) <$> uniqueTransitions numTransitions' maxT
-  weightedTransitions <- (traverse.traverse.traverse) (const genMonoidValue) transitions
+  transitions <-
+    lift
+    $   map (IndexedTransition.fromIndex wtaSpec)
+    <$> uniqueTransitions numTransitions' maxT
+  weightedTransitions <- (traverse . traverse . traverse)
+    (const genMonoidValue)
+    transitions
 
-  let byState = foldl' (\m (State s, t) -> M.insertWith (++) s [t] m) M.empty weightedTransitions
+  let byState = foldl' (\m (State s, t) -> M.insertWith (++) s [t] m)
+                       M.empty
+                       weightedTransitions
 
   let stateVec = V.generate n $ \i -> case M.lookup i byState of
-        Nothing -> V.empty
+        Nothing  -> V.empty
         Just lst -> V.fromList lst
 
   return stateVec
 
 genWTA :: Generator m (WTA m)
 genWTA = asks zeroPolicy >>= \case
-  NumTransitions d -> WTA <$> asks spec <*> genStates <*> (genTransitions' d)
-  ZeroFrequency _ -> WTA <$> asks spec <*> genStates <*> genTransitions
+  NumTransitions d ->
+    WTA <$> asks spec <*> genStates <*> (genTransitionsNumTrans d)
+  ZeroFrequency _ -> WTA <$> asks spec <*> genStates <*> genTransitionsZeroFreq