Switched to unboxed mutable vector

src/Sudoku.hs

@@ -1,8 +1,10 @@
+{-# LANGUAGE LambdaCase #-}
 module Main where
 
 import Control.Applicative ((<|>))
 import Data.Function ((&))
 import qualified Control.Monad
+import qualified Control.Monad.ST
 import qualified Data.Char
 import qualified Data.Function
 import qualified Data.List
@@ -21,6 +23,7 @@ setBits = Data.List.foldl' (Data.Bits..|.)
 
 type Cell = Data.Word.Word16
 type Grid = Data.Vector.Unboxed.Vector Cell
+type MGrid s = Data.Vector.Unboxed.Mutable.STVector s Cell
 type CellIxs = [Int]
 
 isPossible, isFixed :: Cell -> Bool
@@ -46,7 +49,7 @@ readGrid s
 fromXY :: (Int, Int) -> Int
 fromXY (x, y) = x * 9 + y
 
-allRowIxs, allColIxs, allSubGridIxs :: [CellIxs]
+allRowIxs, allColIxs, allSubGridIxs, allIxs :: [CellIxs]
 allRowIxs = [getRow i | i <- [0..8]]
   where getRow n = [ fromXY (n, i) | i <- [0..8] ]
 
@@ -57,8 +60,10 @@ allSubGridIxs = [getSubGrid i | i <- [0..8]]
   where getSubGrid n = let (r, c) = (n `quot` 3, n `mod` 3)
                        in [ fromXY (3 * r + i, 3 * c + j) | i <- [0..2], j <- [0..2] ]
 
+allIxs = concat [allRowIxs, allColIxs, allSubGridIxs]
+
 replaceCell :: Int -> Cell -> Grid -> Grid
-replaceCell i c = Data.Vector.Unboxed.modify (\v -> Data.Vector.Unboxed.Mutable.write v i c)
+replaceCell i c = Data.Vector.Unboxed.modify (\v -> Data.Vector.Unboxed.Mutable.unsafeWrite v i c)
 
 showGrid :: Grid -> String
 showGrid grid = unlines . map (unwords . map (showCell . (grid !))) $ allRowIxs
@@ -97,10 +102,9 @@ exPosAccToList :: ExPosAcc -> [(Int, [Int])]
 exPosAccToList (ExPosAcc v1 v2 v3 v4 v5 v6 v7 v8 v9) =
   [(1, v1), (2, v2), (3, v3), (4, v4), (5, v5), (6, v6), (7, v7), (8, v8), (9, v9)]
 
-exclusivePossibilities :: Grid -> CellIxs -> [Data.Word.Word16]
-exclusivePossibilities grid cellIxs =
-  cellIxs
-  & map (grid !)
+exclusivePossibilities :: [Cell] -> [Data.Word.Word16]
+exclusivePossibilities cells =
+  cells
   & zip ([1..9] :: [Int])
   & filter (isPossible . snd)
   & Data.List.foldl'
@@ -125,41 +129,74 @@ makeCell ys
   | Data.Bits.popCount (makeFixed ys) == 1 = Just $ makeFixed ys
   | otherwise                              = Just $ makePossible ys
 
-pruneCells :: Grid -> CellIxs -> Maybe Grid
-pruneCells grid cellIxs = Control.Monad.foldM pruneCell grid cellIxs
+pruneCells :: MGrid s -> CellIxs -> Control.Monad.ST.ST s (Maybe Bool)
+pruneCells grid cellIxs = do
+  cells <- traverse (Data.Vector.Unboxed.Mutable.unsafeRead grid) cellIxs
+  pruneCells' grid cells cellIxs
+
+pruneCells' :: MGrid s -> [Cell] -> CellIxs -> Control.Monad.ST.ST s (Maybe Bool)
+pruneCells' grid cells = Control.Monad.foldM pruneCell' (Just False)
   where
-    exclusives    = exclusivePossibilities grid cellIxs
+    pruneCell' Nothing _ = return Nothing
+    pruneCell' (Just changed) i = pruneCell i >>= \case
+      Nothing -> return Nothing
+      Just changed' -> return $ Just (changed || changed')
+
+    exclusives    = exclusivePossibilities cells
     allExclusives = setBits Data.Bits.zeroBits exclusives
-    fixeds        = setBits Data.Bits.zeroBits . filter isFixed . map (grid !) $ cellIxs
+    fixeds        = setBits Data.Bits.zeroBits . filter isFixed $ cells
 
-    pruneCell g i =
-      pruneCellByFixed g (i, g ! i) >>= \g' -> pruneCellByExclusives g' (i, g' ! i)
+    pruneCell i = do
+      cell <- Data.Vector.Unboxed.Mutable.unsafeRead grid i
+      pruneCellByFixed (i, cell) >>= \case
+        Nothing -> return Nothing
+        Just changed -> do
+          cell' <- Data.Vector.Unboxed.Mutable.unsafeRead grid i
+          pruneCellByExclusives (i, cell') >>= \case
+            Nothing       -> return Nothing
+            Just changed' -> return $ Just (changed || changed')
 
-    pruneCellByFixed g (i, xs)
-      | isFixed xs = Just g
-      | xs' == xs  = Just g
-      | otherwise  = flip (replaceCell i) g <$> makeCell xs'
+    pruneCellByFixed (i, xs)
+      | isFixed xs = return $ Just False
+      | xs' == xs  = return $ Just False
+      | otherwise  = case makeCell xs' of
+          Nothing -> return Nothing
+          Just cell' -> do
+            Data.Vector.Unboxed.Mutable.unsafeWrite grid i cell'
+            return $ Just True
       where
         xs' = xs Data.Bits..&. Data.Bits.complement fixeds
 
-    pruneCellByExclusives g (i, xs)
-      | isFixed xs                       = Just g
-      | null exclusives                  = Just g
-      | makePossible intersection == xs  = Just g
-      | intersection `elem` exclusives   =
-          flip (replaceCell i) g <$> makeCell intersection
-      | otherwise                      = Just g
+    pruneCellByExclusives (i, xs)
+      | isFixed xs                       = return $ Just False
+      | null exclusives                  = return $ Just False
+      | makePossible intersection == xs  = return $ Just False
+      | intersection `elem` exclusives   = case makeCell intersection of
+          Nothing -> return Nothing
+          Just cell' -> do
+            Data.Vector.Unboxed.Mutable.unsafeWrite grid i cell'
+            return $ Just True
+      | otherwise                        = return $ Just False
       where
         intersection = xs Data.Bits..&. allExclusives
 
-pruneGrid' :: Grid -> Maybe Grid
-pruneGrid' grid =
-  Control.Monad.foldM pruneCells grid allRowIxs
-  >>= flip (Control.Monad.foldM pruneCells) allColIxs
-  >>= flip (Control.Monad.foldM pruneCells) allSubGridIxs
+pruneGrid' :: MGrid s -> Control.Monad.ST.ST s (Maybe Bool)
+pruneGrid' grid = Control.Monad.foldM pruneCells' (Just False) allIxs
+  where
+    pruneCells' Nothing _ = return Nothing
+    pruneCells' (Just changed) ixs = pruneCells grid ixs >>= \case
+      Nothing       -> return Nothing
+      Just changed' -> return $ Just (changed || changed')
 
 pruneGrid :: Grid -> Maybe Grid
-pruneGrid = fixM pruneGrid'
+pruneGrid grid = Control.Monad.ST.runST $ do
+  mGrid <- Data.Vector.Unboxed.unsafeThaw grid
+  fix mGrid
+  where
+    fix mg = pruneGrid' mg >>= \case
+      Nothing -> return Nothing
+      Just True -> fix mg
+      Just False -> Just <$> Data.Vector.Unboxed.unsafeFreeze mg
 
 isGridFilled :: Grid -> Bool
 isGridFilled = not . Data.Vector.Unboxed.any isPossible