--- /dev/null
+-- Writeup at https://work.njae.me.uk/2024/12/09/advent-of-code-2024-day-9/
+
+import AoC
+-- import Data.List
+import Data.Char
+import Data.Maybe
+import qualified Data.IntSet as S
+import qualified Data.IntMap.Strict as M
+
+type Disk = M.IntMap Int
+type Free = S.IntSet
+
+data Region = Free Int -- size
+ | Used Int Int -- size, fileID
+ deriving (Show, Eq)
+
+type RDisk = [Region]
+
+main :: IO ()
+main =
+ do dataFileName <- getDataFileName
+ text <- readFile dataFileName
+ let diskMap = fmap digitToInt text
+ let rdisk = expand diskMap
+ print $ part1 rdisk
+ print $ part2 rdisk
+
+part1, part2 :: RDisk -> Int
+part1 rdisk = checksum $ fst $ packBlocks $ toBlocks rdisk
+part2 rdisk = checksum $ fst $ toBlocks $ packFiles rdisk
+
+-- expand :: [Int] -> (Disk, Free)
+-- expand diskMap = (disk, free)
+-- where (_, _, _, disk, free) = foldl' expandMapItem (True, 0, 0, M.empty, S.empty) diskMap
+
+-- expandMapItem :: (Bool, Int, Int, Disk, Free) -> Int -> (Bool, Int, Int, Disk, Free)
+-- expandMapItem (True, pos, fileID, disk, free) size =
+-- (False, pos + size, fileID + 1, disk', free)
+-- where
+-- fileExtent = take size $ zip [pos..] (repeat fileID)
+-- file = M.fromList fileExtent
+-- disk' = M.union disk file
+-- expandMapItem (False, pos, fileID, disk, free) size =
+-- (True, pos + size, fileID, disk, free')
+-- where
+-- gap = S.fromList $ take size [pos..]
+-- free' = S.union free gap
+
+expand :: [Int] -> RDisk
+expand diskMap = reverse disk
+ where (_, _, _, disk) = foldl' expandRegion (True, 0, 0, []) diskMap
+
+expandRegion :: (Bool, Int, Int, RDisk) -> Int -> (Bool, Int, Int, RDisk)
+expandRegion (True, pos, fID, disk) size =
+ (False, pos + size, fID + 1, (Used size fID) : disk)
+expandRegion (False, pos, fID, disk) 0 =
+ (True, pos, fID, disk)
+expandRegion (False, pos, fID, disk) size =
+ (True, pos + size, fID, (Free size) : disk)
+
+toBlocks :: RDisk -> (Disk, Free)
+toBlocks rdisk = (disk, free)
+ where (_, disk, free) = foldl' toBlock (0, M.empty, S.empty) rdisk
+
+toBlock :: (Int, Disk, Free) -> Region -> (Int, Disk, Free)
+toBlock (pos, disk, free) (Free size) = (pos + size, disk, free')
+ where
+ gap = S.fromList $ take size [pos..]
+ free' = S.union free gap
+toBlock (pos, disk, free) (Used size fileID) = (pos + size, disk', free)
+ where
+ fileExtent = take size $ zip [pos..] (repeat fileID)
+ file = M.fromList fileExtent
+ disk' = M.union disk file
+
+-- Part 1 packing
+
+packBlocks :: (Disk, Free) -> (Disk, Free)
+packBlocks (disk, free) = head $ dropWhile (not . ispackedBlock) $ iterate packBlocksStep (disk, free)
+
+packBlocksStep :: (Disk, Free) -> (Disk, Free)
+packBlocksStep (disk, free)
+ | ispackedBlock (disk, free) = (disk, free)
+ | otherwise = (M.insert to fID disk1, S.insert from free1)
+ where ((from, fID), disk1) = M.deleteFindMax disk
+ (to, free1) = S.deleteFindMin free
+
+ispackedBlock :: (Disk, Free) -> Bool
+ispackedBlock (disk, free) = dMax < fMin
+ where (dMax, _) = M.findMax disk
+ fMin = S.findMin free
+
+checksum :: Disk -> Int
+checksum disk = sum $ fmap (uncurry (*)) $ M.toAscList disk
+
+-- Part 2 packing
+
+packFiles :: RDisk -> RDisk
+packFiles disk = packBelow maxID disk
+ where maxID = maximum $ fmap fileID disk
+
+packBelow :: Int -> RDisk -> RDisk
+packBelow _ [] = []
+packBelow 0 disk = disk
+packBelow fid disk = packBelow (fid - 1) disk'
+ where disk' = tidy $ packFile fid disk
+
+packFile :: Int -> RDisk -> RDisk
+packFile fid disk
+ | isNothing gap = disk
+ | otherwise = prefix ++ [Used fSize fid, Free (gapSize - fSize)] ++ mid ++ [Free fSize] ++ suffix
+ where
+ (orefixMid, ((Used fSize _) : suffix)) = span ((/= fid) . fileID) disk
+ gap = findFree fSize prefixMid
+ (prefix, Free gapSize, mid) = fromJust gap
+
+findFree :: Int -> RDisk -> Maybe (RDisk, Region, RDisk)
+findFree size disk
+ | null suffix = Nothing
+ | otherwise = Just (prefix, head suffix, tail suffix)
+ where (prefix, suffix) = break ((>= size) . freeSize) disk
+
+fileID :: Region -> Int
+fileID (Used _ fID) = fID
+fileID _ = -1
+
+freeSize :: Region -> Int
+freeSize (Free size) = size
+freeSize _ = 0
+
+tidy :: RDisk -> RDisk
+tidy = foldr tidyRegion []
+
+tidyRegion :: Region -> RDisk -> RDisk
+tidyRegion (Free 0) rdisk = rdisk
+tidyRegion (Free size) ((Free size1) : rdisk) = tidyRegion (Free (size + size1)) rdisk
+tidyRegion region rdisk = region : rdisk
+
+
+showDiskFree :: Disk -> Free -> String
+showDiskFree disk free = showDisk disk ++ "\n" ++ showFree free
+
+showDisk :: Disk -> String
+showDisk disk = [showBlock i | i <- [0..pMax]]
+ where (pMax, _) = M.findMax disk
+ showBlock i = maybe '.' intToDigit $ M.lookup i disk
+
+showFree :: Free -> String
+showFree free = [if S.member i free then '+' else '.' | i <- [0..pMax]]
+ where pMax = S.findMax free
\ No newline at end of file