Add missing stuff from day fifteen

day-15-haskell/src/Lib.hs

@@ -4,6 +4,7 @@ module Lib where
 
 import Control.Monad (join)
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import Data.Maybe (fromMaybe, listToMaybe, maybeToList)
 import Flow
 import Text.ParserCombinators.ReadP
@@ -12,11 +13,26 @@ type IntType = Int
 
 data Vec2 = Vec2 IntType IntType
   deriving (Eq, Ord, Show)
+vFromAligned :: Axis -> IntType -> IntType -> Vec2
+vFromAligned Horizontal x y = Vec2 x y
+vFromAligned Vertical y x = Vec2 x y
 vAdd :: Vec2 -> Vec2 -> Vec2
 vAdd (Vec2 aX aY) (Vec2 bX bY) = Vec2 (aX + bX) (aY + bY)
-data Tile = Wall | Box
+vMult :: Vec2 -> Vec2 -> Vec2
+vMult (Vec2 aX aY) (Vec2 bX bY) = Vec2 (aX * bX) (aY * bY)
+vCoord :: Axis -> Vec2 -> IntType
+vCoord Horizontal (Vec2 x _) = x
+vCoord Vertical (Vec2 _ y) = y
+data Tile = Wall | BoxPart { bpSize :: Vec2, bpRel :: Vec2 }
   deriving (Eq, Ord, Show)
+isRootBox :: Tile -> Bool
+isRootBox BoxPart{ bpRel = Vec2 0 0 } = True
+isRootBox _ = False
 data World = World { wMap :: Map.Map Vec2 Tile, wRobot :: Vec2 }
+data Axis = Horizontal | Vertical
+aOrtho :: Axis -> Axis
+aOrtho Horizontal = Vertical
+aOrtho Vertical = Horizontal
 data Direction = DUp | DRight | DDown | DLeft
   deriving (Eq, Ord, Show)
 dVec :: Direction -> Vec2
@@ -24,6 +40,11 @@ dVec DUp = Vec2 0 (-1)
 dVec DRight = Vec2 1 0
 dVec DDown = Vec2 0 1
 dVec DLeft = Vec2 (-1) 0
+dAxis :: Direction -> Axis
+dAxis DUp = Vertical
+dAxis DDown = Vertical
+dAxis DRight = Horizontal
+dAxis DLeft = Horizontal
 type Instruction = Direction
 
 task1 :: String -> IO IntType
@@ -32,10 +53,31 @@ task1 input =
       finalWorld = followInstructions world instructions
   in do
     putStrLn $ showWorld finalWorld
-    wMap finalWorld |> Map.filter ((==) Box) |> Map.keys |> map gps |> sum |> return
+    wMap finalWorld |> Map.filter isRootBox |> Map.keys |> map gps |> sum |> return
 
 task2 :: String -> IO IntType
-task2 _ = return 2
+task2 input =
+  let (world, instructions) = parseFile input
+      tWorld = transformWorld world
+      finalWorld = followInstructions tWorld instructions
+  in do
+    putStrLn $ showWorld finalWorld
+    wMap finalWorld |> Map.filter isRootBox |> Map.keys |> map gps |> sum |> return
+
+  where 
+    transformWorld :: World -> World
+    transformWorld World {wMap, wRobot = Vec2 rX rY } =
+      let transformedMap = wMap |> Map.assocs |> map transformTile |> join |> Map.fromList
+          transformedRobot = Vec2 (rX * 2) rY
+      in World {wMap = transformedMap, wRobot = transformedRobot }
+
+    transformTile :: (Vec2, Tile) -> [(Vec2, Tile)]
+    transformTile (Vec2 tX tY, BoxPart{ bpSize = Vec2 sX sY, bpRel = Vec2 rX rY }) = [
+      (Vec2 (tX * 2) tY, BoxPart{ bpSize = Vec2 (sX * 2) sY, bpRel = Vec2 (rX * 2) rY}),
+      (Vec2 (tX * 2 + 1) tY, BoxPart{ bpSize = Vec2 (sX * 2) sY, bpRel = Vec2 (rX * 2 + 1) rY}) ]
+    transformTile (Vec2 tX tY, tile) = [
+      (Vec2 (tX * 2) tY, tile),
+      (Vec2 (tX * 2 + 1) tY, tile)]
 
 gps :: Vec2 -> IntType
 gps (Vec2 x y) = 100 * y + x
@@ -46,30 +88,58 @@ followInstructions world (instruction:instructions) =
   followInstructions (followInstruction world instruction) instructions
   where
     followInstruction :: World -> Instruction -> World
-    followInstruction world instruction = fromMaybe world $ tryMoveRobot world $ dVec instruction
+    followInstruction world instruction = fromMaybe world $ tryMoveRobot world instruction
 
-    tryMoveRobot :: World -> Vec2 -> Maybe World
-    tryMoveRobot World { wMap, wRobot } vector =
-      let target = wRobot `vAdd` vector
-          targetMap = World { wMap, wRobot = target }
+    tryMoveRobot :: World -> Direction -> Maybe World
+    tryMoveRobot World { wMap, wRobot } direction =
+      let target = wRobot `vAdd` dVec direction
       in case wMap Map.!? target of
-        Nothing -> return targetMap
+        Nothing -> return World { wMap, wRobot = target }
         Just Wall -> Nothing
-        Just Box -> tryMoveBox targetMap target vector
-
-    tryMoveBox :: World -> Vec2 -> Vec2 -> Maybe World
-    tryMoveBox world@World { wMap } from vector =
-      case wMap Map.!? target of
-        Nothing -> Just $ move world
-        Just Wall -> Nothing
-        Just Box -> fmap move $ tryMoveBox world target vector
+        Just BoxPart{} -> do
+          alsoMove <- tryMoveBoxPart world target direction
+          alsoMoveWithTile <- alsoMove |> Set.fromList |> Set.toList |> map (\p -> (p, wMap Map.! p)) |> return
+          removedWorldMap <- foldr Map.delete wMap alsoMove |> return
+          movedWorldMap <- foldr (\(p, t) w -> Map.insert (vAdd p $ dVec direction) t w) removedWorldMap alsoMoveWithTile |> return
+          return World { wMap = movedWorldMap, wRobot = target }
 
+    tryMoveBoxPart :: World -> Vec2 -> Direction -> Maybe [Vec2]
+    tryMoveBoxPart world@World { wMap } pos@(Vec2 pX pY) direction =
+      let allParts = partPositions
+      in do
+        other <- tryMoveParts $ relevantPositions allParts
+        return $ allParts ++ other
       where
-        target :: Vec2
-        target = from `vAdd` vector
+        partPositions :: [Vec2]
+        partPositions = case wMap Map.!? pos of
+          Just BoxPart { bpSize = Vec2 sX sY, bpRel = Vec2 rX rY } ->
+            let bX = pX - rX
+                bY = pY - rY
+            in [ Vec2 x y | x <- [ bX .. (bX + sX - 1) ], y <- [ bY .. (bY + sY - 1) ] ]
+          _ -> []
+
+        relevantPositions :: [Vec2] -> [Vec2]
+        relevantPositions positions =
+          let mainAxis = dAxis direction
+              mainMult = dVec direction |> vCoord mainAxis
+              orthoAxis = aOrtho mainAxis
+          in positions |> map (\p -> (vCoord orthoAxis p, vCoord mainAxis p))
+            |> Map.fromListWith (\a b -> mainMult * max (mainMult * a) (mainMult * b))
+            |> Map.assocs |> map (\(ortho, main) -> vFromAligned mainAxis main ortho)
+
+        tryMoveParts :: [Vec2] -> Maybe [Vec2]
+        tryMoveParts [] = Just []
+        tryMoveParts (p:ps) =
+          let target = vAdd p $ dVec direction
+          in do
+            other <- tryMoveParts ps
+            case wMap Map.!? target of
+              Nothing -> return other
+              Just Wall -> Nothing
+              Just BoxPart{} -> do
+                recBoxParts <- tryMoveBoxPart world target direction
+                return $ other ++ recBoxParts
 
-        move :: World -> World
-        move World { wMap, wRobot } = World { wMap = wMap |> Map.delete from |> Map.insert target Box, wRobot }
 
 showWorld :: World -> String
 showWorld World { wMap, wRobot } =
@@ -82,7 +152,7 @@ showWorld World { wMap, wRobot } =
       | otherwise = case wMap Map.!? pos of
         Nothing -> '.'
         Just Wall -> '#'
-        Just Box -> 'O'
+        Just BoxPart{} -> 'O'
 
 parseFile :: String -> (World, [Instruction])
 parseFile input = case readP_to_S parse input of
@@ -125,7 +195,7 @@ parseWorld = do
 
       toTile :: Char -> Maybe Tile
       toTile '#' = Just Wall
-      toTile 'O' = Just Box
+      toTile 'O' = Just BoxPart { bpSize = Vec2 1 1, bpRel = Vec2 0 0 }
       toTile _ = Nothing
 
 parseInstructions :: ReadP [Instruction]
@@ -133,7 +203,7 @@ parseInstructions = many1 parseInstruction
   where
     parseInstruction :: ReadP Instruction
     parseInstruction = do
-      munch isNewLine
+      _ <- munch isNewLine
       c <- get
       case c of
         '^' -> return DUp