{-# Language ImportQualifiedPost, DataKinds #-}
module KnotHash (knotHash, tieKnots) where

import Advent (chunks)
import Advent.Permutation (mkPermutation, runPermutation, Permutation)
import Data.Bits (Bits(xor))
import Data.Char (ord)
import Data.Foldable (Foldable(foldl'))
import Data.List (foldl1')
import Data.Word (Word8)

-- | Given a rope size and an input string, compute the resulting hash.
knotHash ::
  String  {- ^ input string -} ->
  Integer {- ^ knot value   -}
knotHash :: String -> Integer
knotHash =
   [Word8] -> Integer
bytesToInteger ([Word8] -> Integer) -> (String -> [Word8]) -> String -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Word8] -> Word8) -> [[Word8]] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map ((Word8 -> Word8 -> Word8) -> [Word8] -> Word8
forall a. HasCallStack => (a -> a -> a) -> [a] -> a
foldl1' Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor) ([[Word8]] -> [Word8])
-> (String -> [[Word8]]) -> String -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   Int -> [Word8] -> [[Word8]]
forall a. Int -> [a] -> [[a]]
chunks Int
16 ([Word8] -> [[Word8]])
-> (String -> [Word8]) -> String -> [[Word8]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Word8]
tieKnots ([Int] -> [Word8]) -> (String -> [Int]) -> String -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Int]] -> [Int]) -> (String -> [[Int]]) -> String -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [[Int]]
forall a. Int -> a -> [a]
replicate Int
64 ([Int] -> [[Int]]) -> (String -> [Int]) -> String -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   ([Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
17, Int
31, Int
73, Int
47, Int
23]) ([Int] -> [Int]) -> (String -> [Int]) -> String -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  (Char -> Int) -> String -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
ord

-- | Convert list of bytes into integer in big-endian order.
bytesToInteger :: [Word8] -> Integer
bytesToInteger :: [Word8] -> Integer
bytesToInteger = (Integer -> Word8 -> Integer) -> Integer -> [Word8] -> Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Integer
acc Word8
x -> Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
0x100 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Word8 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
x) Integer
0

-- | Create a rope, tie knots of the given lengths while skipping
-- according to the increasing skip rule.
tieKnots ::
  [Int]   {- ^ knot lengths   -} ->
  [Word8] {- ^ resulting rope -}
tieKnots :: [Int] -> [Word8]
tieKnots [Int]
lengths = (Int -> Word8) -> Permutation 256 -> [Word8]
forall a (n :: Nat). (Int -> a) -> Permutation n -> [a]
runPermutation Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral
                 (Permutation 256 -> [Word8]) -> Permutation 256 -> [Word8]
forall a b. (a -> b) -> a -> b
$ [Permutation 256] -> Permutation 256
forall a. Monoid a => [a] -> a
mconcat [ Int -> Int -> Permutation 256
p Int
o Int
l
                           | (Int
o,Int
l) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int
0..] [Int]
lengths)) [Int]
lengths
                           ]

p :: Int -> Int -> Permutation 256
p :: Int -> Int -> Permutation 256
p Int
o Int
l = (Int -> Int) -> Permutation 256
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation ((Int -> Int) -> Permutation 256)
-> (Int -> Int) -> Permutation 256
forall a b. (a -> b) -> a -> b
$ \Int
i -> if (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
o)Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod`Int
256 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l
                                   then Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
oInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
o
                                   else Int
i