Day 9 - Part 1 refactor with state

src/Day9.hs

@@ -19,7 +19,8 @@ module Day9
 where
 
 import Control.Arrow ((>>>))
-import Data.List (nub)
+import Control.Monad.Trans.State.Strict (State, runState, modify)
+import qualified Data.Set as S
 
 newtype Move = Move (Int, Int)
   deriving (Show, Eq)
@@ -49,17 +50,26 @@ type Pos = (Int, Int)
 
 data Grid = Grid {h :: Pos, t :: Pos} deriving (Show, Eq)
 
-recordPositions :: [Step] -> [Pos]
-recordPositions steps = nub (lastTailPos:positions)
+data MovementLog = MovementLog { visited :: S.Set Pos, recordedSteps :: [Step] }
+  deriving (Show, Eq)
+
+appendOne :: (Pos, Step) -> MovementLog -> MovementLog
+appendOne (p, s) (MovementLog v steps) =
+   MovementLog (S.insert p v) (s:steps)
+
+recordPositions :: [Step] -> ([Pos], [Step])
+recordPositions steps = (allPositions, allSteps)
   where
-    (positions, _, Grid { t = lastTailPos}) = recordPositions' [] steps startGrid
-    recordPositions' :: [Pos] -> [Step] -> Grid -> ([Pos], [Step], Grid)
-    recordPositions' pos [] grid = (pos, [], grid)
-    recordPositions' pos (m : ms) Grid {h = headPos, t = tailPos} =
+    allPositions = S.toList $ S.insert lastTailPos (visited state)
+    allSteps = reverse (recordedSteps state)
+    ((_, Grid { t = lastTailPos}), state) = runState (recordPositions' steps startGrid) (MovementLog S.empty [])
+    recordPositions' :: [Step] -> Grid -> State MovementLog ([Step], Grid)
+    recordPositions' [] grid = return ([], grid)
+    recordPositions' (m : ms) Grid {h = headPos, t = tailPos} =
       let newHead = headPos `step` m
           followStep = follow newHead tailPos
           newTail = tailPos `step` followStep
-       in recordPositions' (tailPos : pos) ms Grid {h = newHead, t = newTail}
+       in modify (appendOne (newTail, followStep)) >> recordPositions' ms Grid {h = newHead, t = newTail}
 
 startGrid :: Grid
 startGrid = Grid {h = p, t = p}
@@ -106,5 +116,5 @@ day9 = do
    input <- readFile "ressources/day09-input"
    putStrLn "Day9"
    let movements = parseMovements input
-   let positions = concatMap normalizeMovement >>> recordPositions >>> length $ movements
+   let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements
    putStrLn ("Number of distinct positions " ++ show positions)

test/Day9Spec.hs

@@ -67,13 +67,13 @@ spec =
         let positions = parseMovements
                 >>> concatMap normalizeMovement
                 >>> recordPositions
+                >>> fst
                 >>> sort
                 $ testInput
         length positions `shouldBe` 13
         positions `shouldBe` expectedPositions
       it "solves the riddle" $ do
         input <- readFile "ressources/day09-input"
-        putStrLn "Day9"
         let movements = parseMovements input
-        let positions = concatMap normalizeMovement >>> recordPositions >>> length $ movements
+        let positions = concatMap normalizeMovement >>> recordPositions >>> fst >>> length $ movements
         positions `shouldBe` 5878