src/advent16/advent16.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5
   6 import Prelude hiding ((++))
   7 import Data.Text (Text)
   8 import qualified Data.Text as T
   9 import qualified Data.Text.IO as TIO
  10
  11 import Text.Megaparsec hiding (State)
  12 import qualified Text.Megaparsec.Lexer as L
  13 import Text.Megaparsec.Text (Parser)
  14 import qualified Control.Applicative as CA
  15
  16 import Control.Monad.State.Lazy
  17 import Control.Monad.Reader
  18
  19 import Data.Vector.Unboxed ((!), (++), (//))
  20 import qualified Data.Vector.Unboxed as V
  21
  22 import qualified Data.IntMap as M
  23
  24
  25 data Step =   Spin Int
  26             | Exchange Int Int
  27             | Partner Char Char
  28             deriving (Show, Eq)
  29
  30 type Dancers = V.Vector Char
  31
  32 type DanceHistory = M.IntMap Dancers
  33
  34 type HistoryRecorder = ReaderT [Step] (State DanceHistory) DanceHistory
  35
  36
  37 startingDancers :: Dancers
  38 startingDancers = V.fromList ['a'..'p']
  39
  40 emptyHistory :: DanceHistory
  41 emptyHistory = M.singleton 0 startingDancers
  42
  43
  44 main :: IO ()
  45 main = do
  46         text <- TIO.readFile "data/advent16.txt"
  47         let instrs = successfulParse text
  48         print $ part1 instrs
  49         print $ part2 instrs
  50
  51
  52 part1 :: [Step] -> Dancers
  53 part1 instrs = evalState (runDance instrs) startingDancers
  54
  55 part2 instrs = (M.!) history (1000000000 `rem` M.size history)
  56     where history = evalState (runReaderT (recordDance startingDancers) instrs) emptyHistory
  57
  58
  59 runDance :: [Step] -> State Dancers Dancers
  60 runDance [] = do dancers <- get
  61                  return dancers
  62 runDance (step:steps) =
  63     do dancers <- get
  64        let dancers' = case step of
  65                         Spin n -> spin n dancers
  66                         Exchange a b -> exchange a b dancers
  67                         Partner a b -> partner a b dancers
  68        put dancers'
  69        runDance steps
  70
  71
  72 recordDance :: Dancers -> HistoryRecorder
  73 recordDance dancers =
  74     do
  75         history <- get
  76         instrs <- ask
  77         let dancers' = evalState (runDance instrs) dancers
  78         if dancers' == startingDancers && (not (history == emptyHistory))
  79         then return history
  80         else do
  81 --                 instrs <- ask
  82 --                 let dancers' = evalState (runDance instrs) dancers
  83                 let history' = M.insert (M.size history) dancers' history
  84                 put history'
  85                 recordDance dancers'
  86
  87 spin :: Int -> Dancers -> Dancers
  88 spin n dancers = back ++ front
  89     where (front, back) = V.splitAt n' dancers
  90           n' = V.length dancers - n
  91
  92 exchange :: Int -> Int -> Dancers -> Dancers
  93 exchange a b dancers = dancers // [(a, dancers!b), (b, dancers!a)]
  94
  95 partner :: Char -> Char -> Dancers -> Dancers
  96 partner a b dancers = exchange a' b' dancers
  97     where a' = V.head $ V.elemIndices a dancers
  98           b' = V.head $ V.elemIndices b dancers
  99
 100
 101 sc :: Parser ()
 102 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 103
 104 -- lexeme  = L.lexeme sc
 105
 106 int :: Parser Int
 107 int = read <$> some digitChar
 108
 109 symb = L.symbol sc
 110 comma = char ','
 111 dancer = oneOf ['a'..'p']
 112
 113 stepsP = stepP `sepBy` comma
 114 stepP = (try spinP) <|> (try exchangeP) <|> partnerP
 115
 116 spinP = Spin <$> (symb "s" *> int)
 117 exchangeP = Exchange <$> (symb "x" *> int) <*> (symb "/" *> int)
 118 partnerP = Partner <$> (symb "p" *> dancer) <*> (symb "/" *> dancer)
 119
 120 successfulParse :: Text -> [Step]
 121 successfulParse input =
 122         case parse stepsP "input" input of
 123                 Left  _error -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 124                 Right steps  -> steps