Candidate lists now working properly

Sudoku.cabal

@@ -94,7 +94,9 @@ executable Sudoku
     -- Other library packages from which modules are imported.
     build-depends:
         base ^>=4.18.2.0,
-        Sudoku
+        Sudoku,
+        vector ^>=0.13.1.0,
+        mtl ^>=2.3.1
 
     -- Directories containing source files.
     hs-source-dirs:   app

app/Main.hs

@@ -1,13 +1,63 @@
 module Main where
 
+import Control.Monad.State (evalState)
+import qualified Data.Vector as V
 import Types
 
 main :: IO ()
 main = do
+  -- putStrLn $ evalState (insert 5 (2, 5) >> showCandidates3) dummyState
+
+  -- print $ evalState (getBoxRange 1) dummyState
+  -- print $ evalState (getBoxRange 2) dummyState
+  -- print $ evalState (getBoxRange 3) dummyState
+  -- print $ evalState (getBoxRange 4) dummyState
+  -- print $ evalState (getBoxRange 5) dummyState
+  -- print $ evalState (getBoxRange 6) dummyState
+  -- print $ evalState (getBoxRange 7) dummyState
+  -- print $ evalState (getBoxRange 8) dummyState
+  -- print $ evalState (getBoxRange 9) dummyState
+
+  -- print $ evalState (getBox (1, 4)) dummyState
+
+  let sudoku = read test :: Sudoku
+
   putStrLn "Reading:"
-  print (read test :: Sudoku)
+  print sudoku
   putStrLn "Pretty Printing:"
-  putStrLn . pretty $ (read test :: Sudoku)
+  putStrLn $ pretty sudoku
+
+  putStrLn "Candidates:"
+  putStrLn $ evalState (insertAll sudoku >> showCandidates3) initState
 
 test :: String
 test = "070000043040009610800634900094052000358460020000800530080070091902100005007040802"
+
+{-
+070 000 043
+040 009 610
+800 634 900
+094 052 000
+358 460 020
+000 800 530
+080 070 091
+902 100 005
+007 040 802
+-}
+dummyState :: SudokuState
+dummyState =
+  SudokuState
+    { dimension = 3,
+      grid = V.replicate 81 0,
+      idxs = V.fromList [(x, y) | x <- [1 .. 9], y <- [1 .. 9]],
+      candidates = V.replicate 81 (V.replicate 9 True)
+    }
+
+initState :: SudokuState
+initState =
+  SudokuState
+    { dimension = 0,
+      grid = V.empty,
+      idxs = V.empty,
+      candidates = V.empty
+    }

src/Types.hs

@@ -2,16 +2,20 @@
 
 module Types where
 
-import Control.Monad.State (State)
+import Control.Monad (forM_, when)
+import Control.Monad.State (MonadState (put), State, get, gets, modify)
 import Data.Char (digitToInt, intToDigit)
-import Data.Vector (Vector)
+import Data.Vector (Vector, slice)
+import qualified Data.Vector as V
+import Debug.Trace
 
 -- * Stateful Sudoku Type
 
 data SudokuState = SudokuState
   { dimension :: Int,
     grid :: Vector Int,
-    idxs :: Vector (Int, Int)
+    idxs :: Vector (Int, Int),
+    candidates :: Vector (Vector Bool)
   }
 
 type SudokuMonad = State SudokuState
@@ -19,20 +23,164 @@ type SudokuMonad = State SudokuState
 getField :: (Int, Int) -> SudokuMonad Int
 getField (x, y) = do
   state <- get
-  return $ (x + y * dimension state) ! grid state
+  return $ grid state ! (x + y * dimension state)
 
 getRow :: (Int, Int) -> SudokuMonad Int
-getRow (x, _) = return x
+getRow (_, y) = return y
 
 getCol :: (Int, Int) -> SudokuMonad Int
-getCol (_, y) = return y
+getCol (x, _) = return x
 
 getBox :: (Int, Int) -> SudokuMonad Int
 getBox (x, y) = do
   d <- gets dimension
-  let xPart = x `div` d
-  let yPart = y `div` d
-  return $ xPart + d * yPart
+  let xPart = ((x - 1) `div` d) + 1
+  let yPart = (y - 1) `div` d
+  return $ xPart + (d * yPart)
+
+getRowRange :: Int -> SudokuMonad [Int]
+getRowRange row = do
+  d <- gets dimension
+  return [x + ((row - 1) * d * d) | x <- [1 .. d * d]]
+
+getColRange :: Int -> SudokuMonad [Int]
+getColRange col = do
+  d <- gets dimension
+  return [col + ((x - 1) * d * d) | x <- [1 .. d * d]]
+
+getBoxRange :: Int -> SudokuMonad [Int]
+getBoxRange box = do
+  d <- gets dimension
+  let (x, y) = getXY box
+  flatIdx <- flattenIdx (x, y)
+  return [flatIdx + n + (m * d * d) | m <- [0 .. d - 1], n <- [0 .. d - 1]]
+  where
+    -- \| Maps a box number to its upper left corner
+    getXY :: Int -> (Int, Int)
+    getXY 1 = (1, 1)
+    getXY 2 = (4, 1)
+    getXY 3 = (7, 1)
+    getXY 4 = (1, 4)
+    getXY 5 = (4, 4)
+    getXY 6 = (7, 4)
+    getXY 7 = (1, 7)
+    getXY 8 = (4, 7)
+    getXY 9 = (7, 7)
+    getXY n = error $ "tried to getXY of box number: " ++ show n
+
+flattenIdx :: (Int, Int) -> SudokuMonad Int
+flattenIdx (x, y) = do
+  d <- gets dimension
+  return (x + (d * d * (y - 1)))
+
+updateVector :: Vector a -> Int -> a -> Vector a
+updateVector v idx x = V.update v (V.singleton (idx - 1, x))
+
+dropCandidate :: Int -> Vector Bool -> Vector Bool
+dropCandidate n bs = updateVector bs n False
+
+dropAllCandidates :: Vector Bool -> Vector Bool
+dropAllCandidates = V.map (const False)
+
+(!) :: Vector a -> Int -> a
+v ! n = v V.! (n - 1)
+
+gridInsert :: Int -> (Int, Int) -> SudokuMonad ()
+gridInsert n idx = do
+  flatIdx <- flattenIdx idx
+  g <- gets grid
+  modify (\s -> s {grid = updateVector g flatIdx n})
+
+removeInRow :: Int -> (Int, Int) -> SudokuMonad ()
+removeInRow n idx = do
+  rr <- getRow idx >>= getRowRange
+  forM_ rr remover
+  where
+    remover :: Int -> SudokuMonad ()
+    remover pos = do
+      cs <- gets candidates
+      let cand = cs ! pos
+      let cand' = dropCandidate n cand
+      modify (\s -> s {candidates = updateVector cs pos cand'})
+
+removeInCol :: Int -> (Int, Int) -> SudokuMonad ()
+removeInCol n idx = do
+  cr <- getCol idx >>= getColRange
+  forM_ cr remover
+  where
+    remover :: Int -> SudokuMonad ()
+    remover pos = do
+      cs <- gets candidates
+      let cand = cs ! pos
+      let cand' = dropCandidate n cand
+      modify (\s -> s {candidates = updateVector cs pos cand'})
+
+removeInBox :: Int -> (Int, Int) -> SudokuMonad ()
+removeInBox n idx = do
+  br <- getBox idx >>= getBoxRange
+  forM_ br remover
+  where
+    remover :: Int -> SudokuMonad ()
+    remover pos = do
+      cs <- gets candidates
+      let cand = cs ! pos
+      let cand' = dropCandidate n cand
+      modify (\s -> s {candidates = updateVector cs pos cand'})
+
+insert :: Int -> (Int, Int) -> SudokuMonad ()
+insert n idx = do
+  gridInsert n idx
+  -- traceM $ "inserting " ++ show n
+  removeInRow n idx
+  removeInCol n idx
+  removeInBox n idx
+  -- remove all candidates for this cell
+  flatIdx <- flattenIdx idx
+  cs <- gets candidates
+  let cand = cs ! flatIdx
+  let cand' = dropAllCandidates cand
+  modify (\s -> s {candidates = updateVector cs flatIdx cand'})
+
+insertAll :: Sudoku -> SudokuMonad ()
+insertAll (Sudoku d g) = do
+  let coords = [(x, y) | y <- [1 .. d * d], x <- [1 .. d * d]]
+  let initialState =
+        SudokuState
+          { dimension = d,
+            grid = V.replicate (d * d * d * d) 0,
+            idxs = V.fromList coords,
+            candidates = V.replicate (d * d * d * d) (V.replicate (d * d) True)
+          }
+  put initialState
+  forM_ (zip g coords) (\(n, idx) -> when (n /= 0) (insert n idx))
+
+showCandidates3 :: SudokuMonad String
+showCandidates3 = do
+  cs <- gets candidates
+  let cs' = V.map (V.filter (/= 0) . V.imap (\idx b -> if b then idx + 1 else 0)) cs
+  return $
+    show (slice 0 9 cs')
+      ++ "\n"
+      ++ show (slice 9 9 cs')
+      ++ "\n"
+      ++ show (slice 18 9 cs')
+      ++ "\n"
+      ++ show (slice 27 9 cs')
+      ++ "\n"
+      ++ show (slice 36 9 cs')
+      ++ "\n"
+      ++ show (slice 45 9 cs')
+      ++ "\n"
+      ++ show (slice 54 9 cs')
+      ++ "\n"
+      ++ show (slice 63 9 cs')
+      ++ "\n"
+      ++ show (slice 72 9 cs')
+
+showGrid :: SudokuMonad String
+showGrid = do
+  g <- gets grid
+  return $ show g
 
 -- * Stateless Sudoku Type
 
@@ -81,17 +229,17 @@ instance Pretty Sudoku where
       prettyBlock (True, n, _) | n == 0 = ""
       prettyBlock (False, n, _) | n == 0 = "|"
       prettyBlock (b, n, x : xs) = intToDigit x : ' ' : prettyBlock (b, n - 1, xs)
-      -- prettyRow (n, dim, list)
+      -- prettyRow (n, d, list)
       prettyRow :: (Int, Int, [Int]) -> String
-      prettyRow (n, dim, xs) | n == 0 = prettyBlock (True, dim, xs)
-      prettyRow (n, dim, xs) = prettyBlock (False, dim, xs) ++ (' ' : prettyRow (n - 1, dim, drop dim xs))
-      -- prettyCol (isLast, n, dim, list)
+      prettyRow (n, d, xs) | n == 0 = prettyBlock (True, d, xs)
+      prettyRow (n, d, xs) = prettyBlock (False, d, xs) ++ (' ' : prettyRow (n - 1, d, drop d xs))
+      -- prettyCol (isLast, n, d, list)
       prettyCol :: (Bool, Int, Int, [Int]) -> String
       prettyCol (True, n, _, _) | n == 0 = ""
-      prettyCol (False, n, dim, _) | n == 0 = replicate (2 * dim * dim - 1 + 2 * (dim - 1)) '-'
+      prettyCol (False, n, d, _) | n == 0 = replicate (2 * d * d - 1 + 2 * (d - 1)) '-'
       prettyCol (_, _, _, []) = ""
-      prettyCol (b, n, dim, xs) = prettyRow (dim - 1, dim, xs) ++ ('\n' : prettyCol (b, n - 1, dim, drop (dim * dim) xs))
-      -- prettyHelper (n, dim, list)
+      prettyCol (b, n, d, xs) = prettyRow (d - 1, d, xs) ++ ('\n' : prettyCol (b, n - 1, d, drop (d * d) xs))
+      -- prettyHelper (n, d, list)
       prettyHelper :: (Int, Int, [Int]) -> String
-      prettyHelper (n, dim, xs) | n == 0 = prettyCol (True, dim, dim, xs)
-      prettyHelper (n, dim, xs) = prettyCol (False, dim, dim, xs) ++ ('\n' : prettyHelper (n - 1, dim, drop (dim * dim) xs))
+      prettyHelper (n, d, xs) | n == 0 = prettyCol (True, d, d, xs)
+      prettyHelper (n, d, xs) = prettyCol (False, d, d, xs) ++ ('\n' : prettyHelper (n - 1, d, drop (d * d * d) xs))