{-# Language QuasiQuotes #-}
{-|
Module      : Main
Description : Day 4 solution
Copyright   : (c) Eric Mertens, 2021
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2015/day/4>

Compute the MD5 hashes of things.

-}
module Main where

import Advent (format)
import Control.Monad                 (replicateM)
import Data.Binary.Get               (runGet, getWord32le)
import Data.Bits                     ((.|.), (.&.), complement, rotateL, xor)
import Data.ByteString.Builder       (Builder, toLazyByteString, lazyByteString, word8, word32LE, word64LE)
import Data.ByteString.Builder.Extra (untrimmedStrategy, toLazyByteStringWith)
import Data.Int                      (Int64)
import Data.List                     (find, foldl')
import Data.Vector                   (Vector)
import Data.Word                     (Word32)
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Char8 as L8
import qualified Data.Vector as Vector

-- |
-- @
-- 282749
-- 9962624
-- @
main :: IO ()
IO ()
main =
  do [Char]
key <- [format|2015 4 %s%n|]
     Maybe Int -> IO ()
forall a. Show a => a -> IO ()
print ([Char] -> Int64 -> Maybe Int
solve [Char]
key Int64
5)
     Maybe Int -> IO ()
forall a. Show a => a -> IO ()
print ([Char] -> Int64 -> Maybe Int
solve [Char]
key Int64
6)

-- | Find the smallest, positive integer that has the specified
-- number of leading zeros in its hex representation.
solve :: String -> Int64 -> Maybe Int
solve :: [Char] -> Int64 -> Maybe Int
solve [Char]
key Int64
n = (Int -> Bool) -> [Int] -> Maybe Int
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Int64 -> ByteString -> Bool
zeros Int64
n (ByteString -> Bool) -> (Int -> ByteString) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Int -> ByteString
adventHash [Char]
key) [Int
1..]

-- | The "advent hash" of a number is the MD5 digest of a key string
-- and a ASCII, base-10 representation of the number.
adventHash ::
  String  {- ^ player key -} ->
  Int     {- ^ number to hash -} ->
  L.ByteString
adventHash :: [Char] -> Int -> ByteString
adventHash [Char]
key Int
i = ByteString -> ByteString
md5 ([Char] -> ByteString
L8.pack ([Char]
key [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i))

-- | Test that the first @n@ digits in hex-representation of
-- the digest are @0@.
zeros :: Int64 -> L.ByteString -> Bool
zeros :: Int64 -> ByteString -> Bool
zeros Int64
n ByteString
bs = (Word8 -> Bool) -> ByteString -> Bool
L.all (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
==Word8
0) (Int64 -> ByteString -> ByteString
L.take Int64
n2 ByteString
bs)
          Bool -> Bool -> Bool
&& (Int64 -> Bool
forall a. Integral a => a -> Bool
even Int64
n Bool -> Bool -> Bool
|| HasCallStack => ByteString -> Int64 -> Word8
ByteString -> Int64 -> Word8
L.index ByteString
bs Int64
n2 Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
0x10)
  where
  n2 :: Int64
n2 = Int64
nInt64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`quot`Int64
2

data Context = Context !Word32 !Word32 !Word32 !Word32

-- > md5 ""
-- d41d8cd98f00b204e9800998ecf8427e
--
-- > md5 "The quick brown fox jumps over the lazy dog."
-- e4d909c290d0fb1ca068ffaddf22cbd0
md5 :: L.ByteString -> L.ByteString
md5 :: ByteString -> ByteString
md5 = Context -> ByteString
finish (Context -> ByteString)
-> (ByteString -> Context) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Context -> Vector Word32 -> Context)
-> Context -> [Vector Word32] -> Context
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Context -> Vector Word32 -> Context
addBlock Context
initialState ([Vector Word32] -> Context)
-> (ByteString -> [Vector Word32]) -> ByteString -> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [Vector Word32]
toBlocks (ByteString -> [Vector Word32])
-> (ByteString -> ByteString) -> ByteString -> [Vector Word32]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
envelope

-- | Extract the final MD5 digest from a context
finish :: Context -> L.ByteString
finish :: Context -> ByteString
finish (Context Word32
a Word32
b Word32
c Word32
d)
  = Int -> Builder -> ByteString
toFixedByteString Int
16
  (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ Word32 -> Builder
word32LE Word32
a Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word32 -> Builder
word32LE Word32
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word32 -> Builder
word32LE Word32
c Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word32 -> Builder
word32LE Word32
d

-- | Pad out an input string to be suitable for breaking into
-- blocks for MD5. This algorithm pads with a @1@ and then
-- as many @0@ bytes as needed so that when the 8-byte length
-- is added that the whole message's length is a multiple of
-- 64-bytes.
envelope :: L.ByteString -> L.ByteString
envelope :: ByteString -> ByteString
envelope ByteString
xs = Builder -> ByteString
toLazyByteString
   (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> Builder
lazyByteString ByteString
xs
  Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
word8          Word8
0x80 -- 0b10000000
  Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
lazyByteString (Int64 -> Word8 -> ByteString
L.replicate Int64
padLen Word8
0)
  Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
word64LE       (Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
bitLen)
  where
  padLen :: Int64
padLen   = (Int64
55 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- ByteString -> Int64
L.length ByteString
xs) Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`mod` Int64
64
  bitLen :: Int64
bitLen   = Int64
8 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
* ByteString -> Int64
L.length ByteString
xs

-- | Break a bytestring with a length that is a multiple of 64
-- into blocks of 16 32-bit words loaded in little-endian order.
toBlocks :: L.ByteString -> [Vector Word32]
toBlocks :: ByteString -> [Vector Word32]
toBlocks
  = (ByteString -> Vector Word32) -> [ByteString] -> [Vector Word32]
forall a b. (a -> b) -> [a] -> [b]
map       ([Word32] -> Vector Word32
forall a. [a] -> Vector a
Vector.fromList ([Word32] -> Vector Word32)
-> (ByteString -> [Word32]) -> ByteString -> Vector Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Get [Word32] -> ByteString -> [Word32]
forall a. Get a -> ByteString -> a
runGet (Int -> Get Word32 -> Get [Word32]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
16 Get Word32
getWord32le))
  ([ByteString] -> [Vector Word32])
-> (ByteString -> [ByteString]) -> ByteString -> [Vector Word32]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> Bool) -> [ByteString] -> [ByteString]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Bool -> Bool
not (Bool -> Bool) -> (ByteString -> Bool) -> ByteString -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bool
L.null)
  ([ByteString] -> [ByteString])
-> (ByteString -> [ByteString]) -> ByteString -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> ByteString) -> ByteString -> [ByteString]
forall a. (a -> a) -> a -> [a]
iterate   (Int64 -> ByteString -> ByteString
L.drop Int64
64)

-- | Point-wise addition of the components of a 'Context'
addState :: Context -> Context -> Context
addState :: Context -> Context -> Context
addState (Context Word32
a Word32
b Word32
c Word32
d) (Context Word32
w Word32
x Word32
y Word32
z) = Word32 -> Word32 -> Word32 -> Word32 -> Context
Context (Word32
aWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+Word32
w) (Word32
bWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+Word32
x) (Word32
cWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+Word32
y) (Word32
dWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+Word32
z)

addBlock ::
  Context ->
  Vector Word32 {- ^ message chunk, 16 elements -} ->
  Context
addBlock :: Context -> Vector Word32 -> Context
addBlock Context
st Vector Word32
m
  = Context -> Context -> Context
addState Context
st
  (Context -> Context) -> Context -> Context
forall a b. (a -> b) -> a -> b
$ Mixer -> [Round] -> Context -> Context
forall {t :: * -> *}.
Foldable t =>
Mixer -> t Round -> Context -> Context
applyRounds Mixer
forall {a}. Bits a => a -> a -> a -> a
m4 [Round]
rs4
  (Context -> Context) -> Context -> Context
forall a b. (a -> b) -> a -> b
$ Mixer -> [Round] -> Context -> Context
forall {t :: * -> *}.
Foldable t =>
Mixer -> t Round -> Context -> Context
applyRounds Mixer
forall {a}. Bits a => a -> a -> a -> a
m3 [Round]
rs3
  (Context -> Context) -> Context -> Context
forall a b. (a -> b) -> a -> b
$ Mixer -> [Round] -> Context -> Context
forall {t :: * -> *}.
Foldable t =>
Mixer -> t Round -> Context -> Context
applyRounds Mixer
forall {a}. Bits a => a -> a -> a -> a
m2 [Round]
rs2
  (Context -> Context) -> Context -> Context
forall a b. (a -> b) -> a -> b
$ Mixer -> [Round] -> Context -> Context
forall {t :: * -> *}.
Foldable t =>
Mixer -> t Round -> Context -> Context
applyRounds Mixer
forall {a}. Bits a => a -> a -> a -> a
m1 [Round]
rs1
  (Context -> Context) -> Context -> Context
forall a b. (a -> b) -> a -> b
$ Context
st
  where
  applyRounds :: Mixer -> t Round -> Context -> Context
applyRounds Mixer
mix t Round
rs Context
st_ = (Context -> Round -> Context) -> Context -> t Round -> Context
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Vector Word32 -> Mixer -> Context -> Round -> Context
doRound Vector Word32
m Mixer
mix) Context
st_ t Round
rs

  m1 :: a -> a -> a -> a
m1 a
b a
c a
d = a
d a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
b a -> a -> a
forall a. Bits a => a -> a -> a
.&. (a
c a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
d))
  m2 :: a -> a -> a -> a
m2 a
b a
c a
d = a
c a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
d a -> a -> a
forall a. Bits a => a -> a -> a
.&. (a
b a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
c))
  m3 :: a -> a -> a -> a
m3 a
b a
c a
d = a
b a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
c a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
d
  m4 :: a -> a -> a -> a
m4 a
b a
c a
d = a
c a -> a -> a
forall a. Bits a => a -> a -> a
`xor` (a
b a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a. Bits a => a -> a
complement a
d)

  rs1 :: [Round]
rs1 = (Int -> Word32 -> Int -> Round)
-> [Int] -> [Word32] -> [Int] -> [Round]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> Word32 -> Int -> Round
Round [Int]
stable1 [Word32]
ktable1 [Int]
gtable1
  rs2 :: [Round]
rs2 = (Int -> Word32 -> Int -> Round)
-> [Int] -> [Word32] -> [Int] -> [Round]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> Word32 -> Int -> Round
Round [Int]
stable2 [Word32]
ktable2 [Int]
gtable2
  rs3 :: [Round]
rs3 = (Int -> Word32 -> Int -> Round)
-> [Int] -> [Word32] -> [Int] -> [Round]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> Word32 -> Int -> Round
Round [Int]
stable3 [Word32]
ktable3 [Int]
gtable3
  rs4 :: [Round]
rs4 = (Int -> Word32 -> Int -> Round)
-> [Int] -> [Word32] -> [Int] -> [Round]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> Word32 -> Int -> Round
Round [Int]
stable4 [Word32]
ktable4 [Int]
gtable4

data Round = Round !Int !Word32 !Int

doRound ::
  Vector Word32 {- ^ message chunk                       -} ->
  Mixer         {- ^ mixing function for this round      -} ->
  Context       {- ^ incoming state                      -} ->
  Round         {- ^ rotation, magic, chunk index -} ->
  Context
doRound :: Vector Word32 -> Mixer -> Context -> Round -> Context
doRound Vector Word32
m Mixer
mixer (Context Word32
a Word32
b Word32
c Word32
d) (Round Int
s Word32
k Int
g) = Word32 -> Word32 -> Word32 -> Word32 -> Context
Context Word32
d (Word32
b Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
z) Word32
b Word32
c
  where
  f :: Word32
f = Mixer
mixer Word32
b Word32
c Word32
d
  y :: Word32
y = Word32
a Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
f Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
k Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Vector Word32
m Vector Word32 -> Int -> Word32
forall a. Vector a -> Int -> a
Vector.! Int
g
  z :: Word32
z = Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
rotateL Word32
y Int
s

type Mixer = Word32 -> Word32 -> Word32 -> Word32

toFixedByteString :: Int -> Builder -> L.ByteString
toFixedByteString :: Int -> Builder -> ByteString
toFixedByteString Int
n = AllocationStrategy -> ByteString -> Builder -> ByteString
toLazyByteStringWith (Int -> Int -> AllocationStrategy
untrimmedStrategy Int
n Int
0) ByteString
L.empty

------------------------------------------------------------------------
-- Magic numbers
------------------------------------------------------------------------

stable1, stable2, stable3, stable4 :: [Int]
stable1 :: [Int]
stable1 = [  Int
7, Int
12, Int
17, Int
22,  Int
7, Int
12, Int
17, Int
22,  Int
7, Int
12, Int
17, Int
22,  Int
7, Int
12, Int
17, Int
22]
stable2 :: [Int]
stable2 = [  Int
5,  Int
9, Int
14, Int
20,  Int
5,  Int
9, Int
14, Int
20,  Int
5,  Int
9, Int
14, Int
20,  Int
5,  Int
9, Int
14, Int
20]
stable3 :: [Int]
stable3 = [  Int
4, Int
11, Int
16, Int
23,  Int
4, Int
11, Int
16, Int
23,  Int
4, Int
11, Int
16, Int
23,  Int
4, Int
11, Int
16, Int
23]
stable4 :: [Int]
stable4 = [  Int
6, Int
10, Int
15, Int
21,  Int
6, Int
10, Int
15, Int
21,  Int
6, Int
10, Int
15, Int
21,  Int
6, Int
10, Int
15, Int
21]


ktable1, ktable2, ktable3, ktable4 :: [Word32]
ktable1 :: [Word32]
ktable1 = [ Word32
0xd76aa478, Word32
0xe8c7b756, Word32
0x242070db, Word32
0xc1bdceee
          , Word32
0xf57c0faf, Word32
0x4787c62a, Word32
0xa8304613, Word32
0xfd469501
          , Word32
0x698098d8, Word32
0x8b44f7af, Word32
0xffff5bb1, Word32
0x895cd7be
          , Word32
0x6b901122, Word32
0xfd987193, Word32
0xa679438e, Word32
0x49b40821]
ktable2 :: [Word32]
ktable2 = [ Word32
0xf61e2562, Word32
0xc040b340, Word32
0x265e5a51, Word32
0xe9b6c7aa
          , Word32
0xd62f105d, Word32
0x02441453, Word32
0xd8a1e681, Word32
0xe7d3fbc8
          , Word32
0x21e1cde6, Word32
0xc33707d6, Word32
0xf4d50d87, Word32
0x455a14ed
          , Word32
0xa9e3e905, Word32
0xfcefa3f8, Word32
0x676f02d9, Word32
0x8d2a4c8a]
ktable3 :: [Word32]
ktable3 = [ Word32
0xfffa3942, Word32
0x8771f681, Word32
0x6d9d6122, Word32
0xfde5380c
          , Word32
0xa4beea44, Word32
0x4bdecfa9, Word32
0xf6bb4b60, Word32
0xbebfbc70
          , Word32
0x289b7ec6, Word32
0xeaa127fa, Word32
0xd4ef3085, Word32
0x04881d05
          , Word32
0xd9d4d039, Word32
0xe6db99e5, Word32
0x1fa27cf8, Word32
0xc4ac5665]
ktable4 :: [Word32]
ktable4 = [ Word32
0xf4292244, Word32
0x432aff97, Word32
0xab9423a7, Word32
0xfc93a039
          , Word32
0x655b59c3, Word32
0x8f0ccc92, Word32
0xffeff47d, Word32
0x85845dd1
          , Word32
0x6fa87e4f, Word32
0xfe2ce6e0, Word32
0xa3014314, Word32
0x4e0811a1
          , Word32
0xf7537e82, Word32
0xbd3af235, Word32
0x2ad7d2bb, Word32
0xeb86d391
          ]

gtable1, gtable2, gtable3, gtable4 :: [Int]
gtable1 :: [Int]
gtable1 = [  Int
0,  Int
1,  Int
2,  Int
3,  Int
4,  Int
5,  Int
6,  Int
7,  Int
8,  Int
9, Int
10, Int
11, Int
12, Int
13, Int
14, Int
15]
gtable2 :: [Int]
gtable2 = [  Int
1,  Int
6, Int
11,  Int
0,  Int
5, Int
10, Int
15,  Int
4,  Int
9, Int
14,  Int
3,  Int
8, Int
13,  Int
2,  Int
7, Int
12]
gtable3 :: [Int]
gtable3 = [  Int
5,  Int
8, Int
11, Int
14,  Int
1,  Int
4,  Int
7, Int
10, Int
13,  Int
0,  Int
3,  Int
6,  Int
9, Int
12, Int
15,  Int
2]
gtable4 :: [Int]
gtable4 = [  Int
0,  Int
7, Int
14,  Int
5, Int
12,  Int
3, Int
10,  Int
1,  Int
8, Int
15,  Int
6, Int
13,  Int
4, Int
11,  Int
2,  Int
9]

initialState :: Context
initialState :: Context
initialState = Word32 -> Word32 -> Word32 -> Word32 -> Context
Context Word32
0x67452301 Word32
0xefcdab89 Word32
0x98badcfe Word32
0x10325476