
   0 | module Days.Day9
   1 | 
   2 | import Data.String
   3 | import Data.Vect
   4 | import Data.SnocList
   5 | import Data.IOArray
   6 | import Data.Fuel
   7 | import Debug.Trace
   8 | import System
   9 | 
  10 | import Control.MonadRec
  11 | 
  12 | import Util
  13 | 
  14 | data Block : Type where
  15 |   Empty : Block
  16 |   File : (id : Nat) -> Block
  17 | %name Block block, block1, block2
  18 | 
  19 | Eq Block where
  20 |   (==) Empty Empty = True
  21 |   (==) Empty (File id) = False
  22 |   (==) (File id) Empty = False
  23 |   (==) (File id) (File id') = id == id'
  24 | 
  25 | Ord Block where
  26 |   compare Empty Empty = EQ
  27 |   compare Empty (File id) = LT
  28 |   compare (File id) Empty = GT
  29 |   compare (File id) (File id') = compare id id'
  30 | 
  31 | showBlock : (largest_id : Nat) -> Block -> String
  32 | showBlock largest_id block =
  33 |   let pad_len : Nat = cast . strLength . show $ largest_id
  34 |       block_str =
  35 |         case block of
  36 |           Empty => "."
  37 |           (File id) => show id
  38 |   in padRight pad_len ' ' block_str
  39 | 
  40 | record Input where
  41 |   constructor MkInput
  42 |   len : Nat
  43 |   block_map : Vect (S len) Block
  44 | %name Input input
  45 | 
  46 | largestId : Input -> Nat
  47 | largestId (MkInput len block_map) =
  48 |   case foldl max Empty block_map of
  49 |     Empty => 0
  50 |     (File id) => id
  51 | 
  52 | Show Input where
  53 |   show input =
  54 |     let largest_id = largestId input
  55 |     in joinBy "" . toList . map (showBlock largest_id) $ input.block_map
  56 | 
  57 | parseInput : String -> Maybe Input
  58 | parseInput str = do
  59 |   xs <- traverse (parsePositive . singleton) . unpack . trim $ str
  60 |   let (len ** block_map) = listToVect $ parseHead xs 0 [<]
  61 |   case len of
  62 |     0 => Nothing
  63 |     (S len) => Just $ MkInput len block_map
  64 |   where
  65 |     addRepeat : (count : Nat) -> (val : Block) -> (acc : SnocList Block) -> SnocList Block
  66 |     addRepeat 0 val acc = acc
  67 |     addRepeat (S k) val acc = addRepeat k val (acc :< val)
  68 |     parseHead : (input : List Nat) -> (idx : Nat) -> (acc : SnocList Block) -> List Block
  69 |     parseHead [] idx acc = toList acc
  70 |     parseHead (x :: []) idx acc =
  71 |       let acc = addRepeat x (File idx) acc
  72 |       in parseHead [] idx acc
  73 |     parseHead (x :: (y :: xs)) idx acc =
  74 |       let acc = addRepeat x (File idx) acc
  75 |           acc = addRepeat y Empty acc
  76 |       in parseHead xs (S idx) acc
  77 | 
  78 | compact : Input -> IO Input
  79 | compact (MkInput len block_map) = do
  80 |   arr <- fromList . map Just . toList $ block_map
  81 |   trSized forever (arr, 0, (max arr) - 1) stepFn
  82 |   xs <- toList arr
  83 |   Just xs <- pure $ traverse id xs
  84 |     | _ => exitFailure
  85 |   Just xs <- pure $ toVect (S len) xs
  86 |     | _ => exitFailure
  87 |   pure $ MkInput len xs
  88 |   where
  89 |     nextEmpty : IOArray Block -> Int -> IO Int
  90 |     nextEmpty xs i =
  91 |       if i >= max xs
  92 |         then pure 0
  93 |         else do
  94 |           Just x <- readArray xs i
  95 |             | _ => exitFailure
  96 |           case x of
  97 |             Empty => pure i
  98 |             File id => nextEmpty xs (i + 1)
  99 |     --
 100 |     prevFile : IOArray Block -> Int -> IO Int
 101 |     prevFile xs i =
 102 |       if i < 0
 103 |         then pure 0
 104 |         else do
 105 |           Just x <- readArray xs i
 106 |             | _ => exitFailure
 107 |           case x of
 108 |             Empty => prevFile xs (i - 1)
 109 |             File id => pure i
 110 |     --
 111 |     swap : IOArray Block -> (i, j : Int) -> IO ()
 112 |     swap xs i j = do
 113 |       Just i_val <- readArray xs i
 114 |         | _ => exitFailure
 115 |       Just j_val <- readArray xs j
 116 |         | _ => exitFailure
 117 |       _ <- writeArray xs i j_val
 118 |       _ <- writeArray xs j i_val
 119 |       pure ()
 120 |     --
 121 |     CState : Type
 122 |     CState = (IOArray Block, Int, Int)
 123 |     --
 124 |     stepFn : (seed2 : Fuel) -> CState -> IO (Step Smaller seed2 CState ())
 125 |     stepFn Dry x = pure $ Done ()
 126 |     stepFn (More y) (xs, empty, file) = do
 127 |       next_empty <- nextEmpty xs empty
 128 |       prev_file <- prevFile xs file
 129 |       if next_empty >= prev_file
 130 |         then pure $ Done ()
 131 |         else do
 132 |           swap xs next_empty prev_file
 133 |           -- TODO: There has got to be a non-hacky way to do this that works with infinite fuel
 134 |           pure $ Cont y (believe_me ()) (xs, next_empty, prev_file)
 135 | 
 136 | checksum : Input -> Nat
 137 | checksum (MkInput len block_map) = checksum' 0 block_map 0
 138 |   where
 139 |     checksum' : (idx : Nat) -> (input : Vect n Block) -> (acc : Nat) -> Nat
 140 |     checksum' idx [] acc = acc
 141 |     checksum' idx (Empty :: xs) acc = checksum' (S idx) xs acc
 142 |     checksum' idx ((File id) :: xs) acc =
 143 |       checksum' (S idx) xs (id * idx + acc)
 144 | 
 145 | export
 146 | part1 : String -> IO (String, Input)
 147 | part1 str = do
 148 |   Just input <- pure $ parseInput str
 149 |     | _ => do
 150 |       putStrLn "Failed to parse input"
 151 |       exitFailure
 152 |   compacted <- compact input
 153 |   pure (show $ checksum compacted, input)
 154 | 
 155 | export
 156 | part2 : Input -> String -> IO String
 157 |