add optparse-generic and better type lookupClosest

spell-checker.cabal

@@ -30,12 +30,13 @@ executable spell-checker
     UndecidableInstances
 
   build-depends:
+    , MemoTrie
     , base
     , containers
-    , MemoTrie
+    , optparse-generic
     , pqueue
-    , vector
     , text
+    , vector
 
   hs-source-dirs:     src
   default-language:   GHC2021

src/Main.hs

@@ -5,15 +5,22 @@ import Data.Map.Strict qualified as M
 import Data.Ord qualified as O
 import Data.Text qualified as T
 import Data.Text.IO qualified as TIO
+import GHC.Generics qualified as G
+import Options.Generic qualified as O
 import PrefixTree qualified as PT
 
-readDict :: IO [T.Text]
-readDict = T.lines <$> TIO.readFile "/usr/share/dict/words"
+data CLIOptions = CLIOptions
+  { word :: T.Text,
+    wordsPath :: FilePath
+  }
+  deriving (Eq, Ord, G.Generic, O.ParseRecord)
+
+readDict :: FilePath -> IO [T.Text]
+readDict = fmap T.lines . TIO.readFile
 
 main :: IO ()
 main = do
-  allWords <- readDict
-  let word = "bril"
+  CLIOptions {..} <- O.getRecord "dictionary options"
+  allWords <- readDict wordsPath
   let vals = PT.fromList . map T.unpack $ allWords
-  -- print $ M.size . PT.children $ vals
-  print $ PT.lookupClosest word vals
+  print $ PT.lookupClosest (T.unpack word) vals

src/PrefixTree.hs

@@ -17,6 +17,9 @@ data Tree a = (Ord a) =>
   { isTermination :: Bool,
     children :: M.Map a (Tree a)
   }
+deriving instance Eq (Tree a)
+deriving instance Ord (Tree a)
+deriving instance Show a => Show (Tree a)
 
 empty :: (Ord a) => Tree a
 empty = Tree False M.empty
@@ -28,38 +31,55 @@ insert :: Tree a -> [a] -> Tree a
 insert (Tree {..}) [] = Tree {isTermination = True, ..}
 insert (Tree {..}) (x : xs) = Tree isTermination . flip (M.insert x) children . flip insert xs . My.fromMaybe empty . M.lookup x $ children
 
+data Step a = Ord a => Step {
+  rest :: [a],
+  soFar :: [a],
+  tree :: Tree a
+                                     }
+
+data StepLookupOutput a = End [a] | Recurse (Step a)
+
+deriving instance Eq (Step a)
+deriving instance Ord (Step a)
+deriving instance Show a => Show (Step a)
+
 lookup :: (Ord a) => [a] -> Tree a -> Bool
 lookup [] = const True
 lookup (x : xs) = maybe False (PrefixTree.lookup xs) . M.lookup x . children
 
 lookupClosest :: forall a. (Ord a, Show a) => [a] -> Tree a -> Maybe (Word, [a])
-lookupClosest xs tree = lookupClosest' . PQ.singleton 0 . Left $ ([], xs, tree)
+lookupClosest xs tree = lookupClosest' . PQ.singleton 0 . Recurse $ Step {rest = xs, soFar = [], ..}
 
-lookupClosest' :: forall a. (Ord a, Show a) => PQ.MinPQueue Word (Either ([a], [a], Tree a) [a]) -> Maybe (Word, [a])
+lookupClosest' :: forall a. (Ord a, Show a) => PQ.MinPQueue Word (StepLookupOutput a) -> Maybe (Word, [a])
 lookupClosest' pq =
   case PQ.minViewWithKey pq of
     Nothing -> Nothing
-    Just ((cost, Right xs), _) -> Just (cost, reverse xs)
-    Just ((cost, Left (soFar, xs, tree)), pq') ->
+    Just ((cost, End xs), _) -> Just (cost, reverse xs)
+    Just ((cost, (Recurse step)), pq') ->
       lookupClosest'
         . PQ.union pq'
         . PQ.fromList
         . map (Bi.first (+ cost))
-        . lookupClosestStep soFar xs
-        $ tree
+        . lookupClosestStep $ step
 
-lookupClosestStep :: forall a. (Ord a, Show a) => [a] -> [a] -> Tree a -> [(Word, Either ([a], [a], Tree a) [a])]
-lookupClosestStep soFar [] (Tree {..})
-  | isTermination = [(0, Right soFar)]
-  | otherwise = [(8, Right soFar)]
-lookupClosestStep soFar xs (Tree {..}) | M.null children = [(fromIntegral . length $ xs, Right soFar)]
-lookupClosestStep soFar (x : xs) (Tree {..}) = map (Bi.second Left) $ rest ++ perfectMatch
+lookupClosestStep :: forall a. (Ord a, Show a) => Step a -> [(Word, StepLookupOutput a)]
+lookupClosestStep (Step {rest = [], tree = (Tree {..}), ..})
+  | isTermination = [(0, End soFar)]
+  | otherwise = map ((1,) . Recurse .  uncurry (Step []) . Bi.first (: soFar)) . M.assocs $ children
+lookupClosestStep (Step {tree = (Tree {..}), ..}) | M.null children = [(L.genericLength rest, End soFar)]
+lookupClosestStep (Step {tree = (Tree {..}), rest = (x:xs), ..}) = map (Bi.second Recurse) $ rest' ++ perfectMatch
  where
-  rest :: [(Word, ([a], [a], Tree a))]
-  rest = concatMap (\x' ->[charReplacedResult x', charSkippedResult x']) . M.assocs . M.delete x $ children
+  rest' :: [(Word, Step a)]
+  rest' = concatMap (\x' -> [charReplacedResult x', charSkippedResult x']) . M.assocs . M.delete x $ children
 
-  perfectMatch :: [(Word, ([a], [a], Tree a))]
-  perfectMatch = My.maybeToList . fmap ((0,) . (x : soFar,xs,)) . M.lookup x $ children
+  perfectMatch :: [(Word, Step a)]
+  perfectMatch = My.maybeToList . fmap ((0,) .Step xs (x : soFar)) . M.lookup x $ children
 
-  charReplacedResult (skippedChar, tree') = (1, (skippedChar : soFar, xs, tree'))
-  charSkippedResult (skippedChar, tree') =  (1, (skippedChar : soFar, x:xs, tree'))
+  charReplacedResult = missedCharResult xs
+  charSkippedResult = missedCharResult (x : xs)
+
+  missedCharResult newXs (skippedChar, tree') = (1,) $ Step {
+    tree = tree',
+    soFar = skippedChar : soFar,
+    rest = newXs
+                                                     }