rainbow: Solve quickly.

Oh my. There were a *lot* of solutions!

Rainbow.hs

@@ -6,8 +6,6 @@ import Control.Monad.State
 import Data.List
 import Data.List.Split
 import qualified Data.Map as M
-import qualified Data.Set as S
-import Debug.Trace
 
 data Color = None | R | O | Y | G | B | V
     deriving (Enum, Eq, Ord, Show)
@@ -18,7 +16,12 @@ type Coord = (Int, Int)
 type Slot = (Height, Color)
 type Board = M.Map Coord Slot
 
-type Solver = StateT (Board, S.Set Slot) []
+type Solver = StateT (Board, [Slot]) []
+
+-- This particular form comes from Cale, because Cale is awesome.
+withEach :: [a] -> [(a, [a])]
+withEach [] = []
+withEach (a:as) = (a, as) : [(b, a : bs) | (b, bs) <- withEach as]
 
 readPuzzle :: String -> [Height]
 readPuzzle s = do
@@ -30,7 +33,7 @@ heightToSlot :: Height -> Slot
 heightToSlot h = (h, None)
 
 makeBoard :: [Height] -> Board
-makeBoard hs = at (0, 0) . _Just . _2 .~ R $ M.fromList $
+makeBoard hs = M.fromList $
     zip [(i, j) | i <- [0 .. 5], j <- [0 .. 5] ] (map heightToSlot hs)
 
 showBoard :: Board -> String
@@ -41,8 +44,8 @@ showBoard b = unlines ls
         results = map show colors
         colors = map (snd . snd) . sort $ M.toList b
 
-initialState :: S.Set Slot
-initialState = S.fromList [(h, c) | h <- [One .. Six], c <- [R .. V] ]
+initialState :: [Slot]
+initialState = [(h, c) | h <- [One .. Six], c <- [R .. V]]
 
 colorAt :: Coord -> Solver Color
 colorAt c = do
@@ -61,22 +64,16 @@ unique (x, y) = do
         color' <- colorAt (x, y')
         guard $ color /= color'
 
-updateBoard :: Coord -> Slot -> (Board, S.Set Slot) -> (Board, S.Set Slot)
-updateBoard c (height, color) board =
-    board & _1 . at c . _Just . _2 .~ color & _2 %~ S.delete (height, color)
-
-derp x = traceShow x x
-
 assignCoord :: Coord -> Solver ()
 assignCoord c = do
     Just (h, color) <- use $ _1 . at c
-    colors <- uses _2 $ S.map snd . S.filter ((h ==) . fst)
-    guard $ not (S.null colors)
-    board <- use id
-    -- Update each part of the state.
-    let states = map (\color' -> ((), updateBoard c (h, color') board)) $
-                    S.toList colors
-    StateT (\_ -> states)
+    colors <- use _2
+    -- Pull out the different possibilities.
+    ((h', color'), colors') <- lift $ withEach colors
+    -- Guard away the ones that we don't like.
+    guard $ h == h'
+    -- Update the state.
+    _1 . at c . _Just . _2 .= color'
     -- And now ensure that we're not propagating duplicate states.
     unique c
 
@@ -93,4 +90,7 @@ main = do
     print board
     putStr $ showBoard board
     let solved = execStateT solve (board, initialState)
-    mapM_ (putStr . showBoard . fst) solved
+    -- putStr $ "There are " ++ show (length solved) ++ " solutions."
+    forM_ (zip solved [0 ..]) $ \((b, _), i) -> do
+        putStr $ "Solution " ++ show i ++ ":\n"
+        putStr $ showBoard b