advent12/src/advent12.hs

   1 import Data.Text (Text)
   2 import qualified Data.Text.IO as TIO
   3
   4 import Data.Void (Void)
   5
   6 import Text.Megaparsec hiding (State)
   7 import Text.Megaparsec.Char
   8 import qualified Text.Megaparsec.Char.Lexer as L
   9 import qualified Control.Applicative as CA
  10
  11 import Linear (V3(..), (^+^), (^-^))
  12
  13 import qualified Data.Set as S
  14
  15 -- import Data.List (foldl')
  16 -- import Data.Set ((\\))
  17 -- import qualified Data.Map.Strict as M
  18 -- import Data.Map.Strict ((!))
  19
  20
  21 type Vec = V3 Integer
  22 data Planet = Planet { _pos :: Vec, _vel :: Vec} deriving (Show, Eq, Ord)
  23 type Planets = S.Set Planet
  24
  25
  26 main :: IO ()
  27 main = do
  28         text <- TIO.readFile "data/advent12a.txt"
  29         let planetsT = successfulParse text
  30         let planets = enplanet planetsT
  31         print planets
  32         print $ part1 planets
  33
  34
  35 part1 planets = take 12 $ simulate planets
  36
  37
  38 enplanet  = S.fromList . map (\p -> Planet {_pos = p, _vel = (V3 0 0 0)} )
  39
  40 _x (V3 x _ _) = x
  41 _y (V3 _ y _) = y
  42 _z (V3 _ _ z) = z
  43
  44
  45 gravity (V3 x y z) = V3 (signum x) (signum y) (signum z)
  46
  47
  48 simulate = iterate simulationStep
  49
  50 simulationStep planets = planets''
  51     where   planets' = applyGravity planets
  52             planets'' = applyVelocity planets'
  53
  54
  55 applyGravity planets = S.map (applyGravityHere planets) planets
  56
  57 applyGravityHere planets here = S.foldl' updateGravity here planets
  58
  59 updateGravity here there = here { _vel = vel'}
  60     where   vel = _vel here
  61             vel' = vel ^+^ gravity ((_pos there) ^-^ (_pos here))
  62
  63
  64 applyVelocity = S.map applyVelocityHere
  65
  66 applyVelocityHere here = here {_pos = (_pos here) ^+^ (_vel here)}
  67
  68
  69
  70 -- Parse the input file
  71 type Parser = Parsec Void Text
  72
  73 sc :: Parser ()
  74 sc = L.space (skipSome spaceChar) CA.empty CA.empty
  75 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  76
  77 lexeme  = L.lexeme sc
  78 integer = lexeme L.decimal
  79 signedInteger = L.signed sc integer
  80 symb = L.symbol sc
  81 equalP = symb "="
  82 commaP = symb ","
  83 identifierP = some alphaNumChar <* sc
  84 openBracketP = symb "<"
  85 closeBracketP = symb ">"
  86
  87 planetsP = many planetP
  88
  89 planetP = (between openBracketP closeBracketP) coordsP
  90
  91 coordsP = envector <$> (coordP `sepBy` commaP)
  92     where envector [x, y, z] = V3 x y z
  93 coordP = identifierP *> equalP *> signedInteger
  94
  95
  96 successfulParse :: Text -> [Vec]
  97 successfulParse input =
  98         case parse planetsP "input" input of
  99                 Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 100                 Right planets -> planets