Day 9 Part 2

09/Main.idr

@@ -23,12 +23,6 @@ parseDirection input =
     "D" => Just Down
     x => Nothing
 
-invert : Direction -> Direction
-invert Up = Down
-invert Right = Left
-invert Left = Right
-invert Down = Up
-
 -- Apply a direction to a pair of Ints
 applyDirection : Direction -> (Int, Int) -> (Int, Int)
 applyDirection Up (x, y) = (x, y + 1)
@@ -65,51 +59,68 @@ breakdown : (motion : Motion) -> Vect (distance motion) Direction
 breakdown (Move direction 0) = []
 breakdown input@(Move direction (S k)) = direction :: breakdown (assert_smaller input (Move direction k))
 
--- The state of a head and tail
-data State: Type where
-  MkState : (head : (Int, Int)) -> (tail : (Int, Int)) -> State
-%name State state, state1, state2
-
-Show State where
-  show (MkState head tail) = "Head: " ++ show head ++ " Tail: " ++ show tail
-
-emptyState : State
-emptyState = MkState (0,0) (0,0)
-
-tail : State -> (Int, Int)
-tail (MkState head x) = x
-
 touching : (Int, Int) -> (Int, Int) -> Bool
 touching (x, y) (z, w) =
   let (u, v) = (abs (x - z), abs (y - w)) in
   u <= 1 && v <= 1
 
--- Move the head in a particular direction, causing the tail to follow if needed
-moveHead : Direction -> State -> State
-moveHead direction (MkState head tail) =
-  let newHead = applyDirection direction head in
-  if touching newHead tail
-    then MkState newHead tail
-    else let newTail = applyDirection (invert direction) newHead in
-         MkState newHead newTail
+data State : Nat -> Type where
+  MkState : (head : (Int, Int)) -> (tails : Vect n (Int, Int)) -> State n
+%name State state, state1, state2
 
--- Apply a movement to the head, returning the sequence of tail positions, including the starting
--- position
-applyMotion : (motion : Motion) -> State -> (State, Vect (S (distance motion)) State)
-applyMotion motion state = helper (breakdown motion) state
-  where helper : (directions : Vect n Direction) -> State -> (State, Vect (S n) State)
+emptyState : {n : Nat} -> State n
+emptyState = MkState (0,0) (replicate _ (0,0))
+
+Show (State n) where
+  show (MkState head tails) = "Head: " ++ show head ++ " Tails: " ++ show tails
+
+stateTails : State n -> Vect n (Int, Int)
+stateTails (MkState head xs) = xs
+
+sign : Int -> Int
+sign i = if i == 0 then 0 else div i (abs i)
+
+mapP : (a -> b) -> (a, a) -> (b, b)
+mapP f (x, y) = (f x, f y)
+
+zipP : (a -> b -> c) -> (a, a) -> (b, b) -> (c, c)
+zipP f (x, y) (z, w) = (f x z, f y w)
+
+-- Increment the tail
+incTail : (head : (Int, Int)) -> (tail : (Int, Int)) -> (Int, Int)
+incTail head tail =
+  let diff = zipP (-) head tail
+      absDiff = mapP abs diff
+      signs = mapP sign diff
+  in if touching head tail
+     then tail
+     else zipP (+) tail signs
+
+incTails : {n: Nat} -> (head : (Int, Int)) -> (tails : Vect n (Int, Int)) -> Vect n (Int, Int)
+incTails head [] = []
+incTails head (x :: xs) =
+  let newTail = incTail head x in
+  newTail :: incTails newTail xs
+
+moveHead' : {n: Nat} -> Direction -> State n -> State n
+moveHead' direction (MkState head tails) =
+  let newHead = applyDirection direction head in
+  MkState newHead (incTails newHead tails)
+
+applyMotion' : {m: Nat} -> (motion : Motion) -> State m -> (State m, Vect (S (distance motion)) (State m))
+applyMotion' motion state = helper (breakdown motion) state
+  where helper : {m: Nat} -> (directions : Vect n Direction) -> State m -> (State m, Vect (S n) (State m))
         helper [] state = (state, [state])
         helper (x :: xs) state =
-          let newState = moveHead x state
+          let newState = moveHead' x state
               (resState, rest) = helper xs newState
           in (resState, state :: rest)
 
--- Apply all the motions in a list to the head, returning the full sequence of positions
-applyMotions : (motions : List Motion) -> State -> (State, List State)
-applyMotions [] state = (state, [state])
-applyMotions (x :: xs) state =
-  let (newState, motions) = applyMotion x state
-      (outputState, rest) = applyMotions xs newState
+applyMotions' : {m : Nat} -> (motions : List Motion) -> State m -> (State m, List (State m))
+applyMotions' [] state = (state, [state])
+applyMotions' (x :: xs) state =
+  let (newState, motions) = applyMotion' x state
+      (outputState, rest) = applyMotions' xs newState
   in (outputState, (toList motions) ++ rest)
 
 simple = """
@@ -123,13 +134,25 @@ simple = """
   R 2
   """
 
-testPart1 : IO ()
-testPart1 =
-  case parseMotions simple of
+complex = """
+  R 5
+  U 8
+  L 8
+  D 3
+  R 17
+  D 10
+  L 25
+  U 20
+  """
+
+testPart : Nat -> (input : String) -> IO ()
+testPart Z input = putStrLn "No Tail!"
+testPart (S m) input =
+  case parseMotions input of
     Nothing => putStrLn "Failed to parse motions"
     Just motions =>
-      let (state, states) = applyMotions motions emptyState
-          tails = map tail states
+      let (state, states) = applyMotions' motions (emptyState {n = (S m)})
+          tails = map (last . stateTails) states
           uniqueTails = nub tails
       in do putStrLn "Movements:"
             traverse_ printLn motions
@@ -138,15 +161,31 @@ testPart1 =
             putStrLn "\nUnique Tails:"
             printLn $ length uniqueTails
 
-part1 : String -> Maybe Nat
-part1 input =
+testPart1 : IO ()
+testPart1 = testPart 1 simple
+
+testPart2 : IO ()
+testPart2 =
+  do testPart 9 simple
+     testPart 9 complex
+
+part : Nat -> String -> Maybe Nat
+part 0 str = ?part_rhs_0
+part (S k) input =
   case parseMotions input of
     Nothing => Nothing
     Just motions =>
-      let (state, states) = applyMotions motions emptyState
-          tails = map tail states
+      let (state, states) = applyMotions' motions (emptyState {n = (S k)})
+          tails = map (last . stateTails) states
       in Just . length . nub $ tails
 
+
+part1 : String -> Maybe Nat
+part1 input = part 1 input
+
+part2 : String -> Maybe Nat
+part2 input = part 9 input
+
 partial main : IO ()
 main =
   do file <- readFile "input"
@@ -156,4 +195,7 @@ main =
         case part1 contents of
           Nothing => printLn "Error in part 1"
           Just part1Count =>
-            putStrLn ("Part 1: " ++ show part1Count)
+            do putStrLn ("Part 1: " ++ show part1Count)
+               case part2 contents of
+                 Nothing => printLn "Error in part 2"
+                 Just part2Count => putStrLn ("Part 2: " ++ show part2Count)