cabal-version:      2.4
name:               keccak
version:            0.1.0.0
synopsis:           Playing with Keccak in Haskell.
-- description:
homepage:           https://nest.pijul.com/sellout/crypto-junk
bug-reports:        https://nest.pijul.com/sellout/crypto-junk/discussions
license:            AGPL-3.0-or-later
author:             Greg Pfeil
maintainer:         greg@technomadic.org
copyright:          2022 Greg Pfeil
category:           Cryptography
extra-source-files: CHANGELOG.md
library
    exposed-modules:  Keccak
                    , Sponge
    ghc-options:     -Weverything
    build-depends:    base
                    , bitvec
                    , fin
                    , lens
                    , vec
                    , vec-lens
                    , vector
                    , yaya
    -- These are required by @-Weverything@.
    default-extensions: DerivingStrategies
                      , ImportQualifiedPost
                      , NoImplicitPrelude
    default-language: Haskell2010

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- Needed for nested type family application in `Exp`.
{-# LANGUAGE UndecidableInstances #-}
-- __FIXME__: These are required because of two things. The first is
--           "Data.Bits". We should replace this with a statically-sized
--            alternative @StaticBits (n :: `Nat`) a@ with indexed lookups,
--            etc. It might still be unsafe for primitive types, though. The
--            second is "Data.Vector.Generic". This is currently an efficient
--            representation of arbitrary-length bit strings. Hopefully we can
--            find an alternative for this, too.
{-# OPTIONS_GHC -Wno-missing-safe-haskell-mode -Wno-unsafe #-}
-- Because of @`Bits` `Vec`@ instances.
{-# OPTIONS_GHC -Wno-orphans #-}
-- | This is an implementation of [Keccak](https://keccak.team/keccak.html) with
--   a lot of static guarantees.
--
--   There are some deviations from the spec:
-- * `keccak_f`, rather than taking the value of @b@ takes the value of @l@, so
--   where the spec says "Keccak-/f/[1600]", we would write
--   @`keccak_f` (Proxy \@(`Nat.fromGHC` 6))@.
-- * `keccak_rc` and `keccak_c` each take an extra `Proxy` (@l@ and @r@,
--   respectively), but these are uniquely determined by the other parameters,
--   so it's a minor annoyance.
module Sponge
  ( sponge,
    Block,
    W,
    WBits,
    State,
    B,
  )
where
import Control.Applicative (Applicative (..))
import Control.Arrow (Arrow (..))
import Control.Category (Category (..))
import Data.Bit (Bit (..), Vector)
import Data.Bitraversable (Bitraversable (..))
import Data.Bits (Bits (..), FiniteBits (..))
import Data.Bool (Bool (..))
import Data.Foldable (Foldable (..))
import Data.Functor (Functor (..), (<$>))
import Data.Maybe (Maybe)
import Data.Proxy (Proxy (..))
import Data.Semigroup (Semigroup (..))
import Data.Tuple (uncurry)
import Data.Type.Nat (Nat (..))
import Data.Type.Nat qualified as Nat
import Data.Type.Nat.LE qualified as Nat
import Data.Vec.Lazy (Vec (..))
import Data.Vec.Lazy qualified as Vec
import Data.Vector.Generic qualified as V
import Yaya.Fold (Corecursive (..), Projectable (..), Steppable (..))
import Yaya.Zoo (Stream)
import Prelude (Num (..), fromIntegral, undefined, ($))
-- | This allows us to do bit operations at whatever level of the arrays that we
--   want.
instance (Nat.SNatI n, Bits a) => Bits (Vec n a) where
  (.&.) = Vec.zipWith (.&.)
  (.|.) = Vec.zipWith (.|.)
  xor = Vec.zipWith xor
  complement = fmap complement
  bitSizeMaybe _ = (Nat.reflectToNum (Proxy @n) *) <$> bitSizeMaybe @a undefined
  {-# INLINEABLE bitSizeMaybe #-}
instance (Nat.SNatI n, FiniteBits a) => FiniteBits (Vec n a) where
  finiteBitSize _ = Nat.reflectToNum (Proxy @n) * finiteBitSize @a undefined
type family Exp (n :: Nat) (m :: Nat) :: Nat where
  Exp 'Z ('S _) = 'Z
  Exp ('S _) 'Z = 'S 'Z
  Exp n ('S 'Z) = n -- avoids an extra @Nat.Mult n 1@ at the end
  Exp n ('S m) = Nat.Mult n (Exp n m)
-- = Keccak-specific definitions
-- | The number of bits in a row or column.
type Block = Nat.FromGHC 5
-- | The number of bits in a lane.
type W l = Exp (Nat.FromGHC 2) l
-- | A structure containing @`W` l@ bits.
--
--  __TODO__: Replace this with the more efficient type family below (once we
--            switch to using bitwise operations on @`Vec n `Bool`@).
type WBits (l :: Nat) = Vec (W l) Bool
-- type family WBits (l :: Nat) where
--   WBits 'Z = Bool
--   -- | This could use better types for the 2- and 4-bit cases.
--   WBits ('S 'Z) = Vec (Nat.FromGHC 2) Bool
--   WBits ('S ('S 'Z)) = Vec (Nat.FromGHC 4) Bool
--   WBits ('S ('S ('S 'Z))) = Word8
--   WBits ('S ('S ('S ('S 'Z)))) = Word16
--   WBits ('S ('S ('S ('S ('S 'Z))))) = Word32
--   WBits ('S ('S ('S ('S ('S ('S 'Z)))))) = Word64
type State l = Vec Block (Vec Block (WBits l))
-- | The width of the permutation.
type B l = Nat.Mult Block (Nat.Mult Block (W l))
restructureVec ::
  (Nat.SNatI x, Nat.SNatI y, Nat.SNatI (W l), Nat.SNatI (Nat.Mult y (W l))) =>
  Proxy l ->
  Vec (Nat.Mult x (Nat.Mult y (W l))) Bool ->
  Vec x (Vec y (WBits l))
restructureVec Proxy = fmap Vec.chunks . Vec.chunks
unstructureVec ::
  Proxy l ->
  Vec x (Vec y (WBits l)) ->
  Vec (Nat.Mult x (Nat.Mult y (W l))) Bool
unstructureVec Proxy = Vec.concat . fmap Vec.concat
data ComposedVec w n a = ComposedVec {getVec :: (Vec n a), getWritten :: w}
-- | Collects exactly @n@ elements from the an infinite stream.
takeFromStream :: Nat.SNatI n => Stream a -> Vec n a
takeFromStream stream =
  getVec $
    Nat.induction1
      (ComposedVec VNil stream)
      (\(ComposedVec v s) -> uncurry ComposedVec . first (::: v) $ project s)
-- |
--
--  __FIXME__: Get rid of this function. Currently it breaks if we just call
--            `takeFromStream`.
takeEnough ::
  Nat.SNatI n => Stream (Vec ('S r) Bool) -> Vec ('S n) (Vec ('S r) Bool)
takeEnough = takeFromStream
-- |
--
--  __NB__: In [/Cryptographic sponge
--          functions/](https://keccak.team/files/CSF-0.1.pdf#page=14), the
--          sponge construction relies on @b@. However, it never includes it as
--          a parameter. Here, because @b = r + c@ and because subtraction is
--          difficult, we replace @b - r@ with @c@.
--
--  __FIXME__: The constraints around @n@ here trigger a cascade of annoyances,
--             where the caller eventually needs to explicitly provide the value
--             satisfying @⌈l' / r⌉@.
--
--  __FIXME__: This returns in `Maybe` because @cut@ can’t tell that the length
--             of the input `Vector` is exactly some multiple of @r@. We should
--             change how padding works to provide this guarantee at the type
--             level. Also, a result of `Nothing` would indicate a bug in this
--             implementation. Which means the failure should at least be more
--             informative, but possibly even an exception. However, returning
--            `Maybe` is a good reminder to fix this, as it makes the
--             shortcoming apparent to all callers.
sponge ::
  forall l' l r c n.
  ( Nat.SNatI n,
    Nat.SNatI (W l),
    Bits (WBits l),
    Nat.SNatI r,
    Nat.SNatI c,
    Nat.SNatI (Nat.Mult Block (W l)),
    B l ~ Nat.Plus ('S r) ('S c),
    Nat.LE ('S r) (B l),
    Nat.LE ('S l') (Nat.Mult ('S n) ('S r))
  ) =>
  Proxy l ->
  (State l -> State l) ->
  (Proxy ('S r) -> Nat -> Vector Bit) ->
  Proxy ('S r) ->
  Proxy ('S c) ->
  -- | The optimal value here is  ⌈l' / r⌉, but any value /at least/ that high
  --   will suffice.
  Proxy ('S n) ->
  Vector Bit ->
  Maybe (Vec ('S l') Bool)
sponge l f pad r Proxy Proxy =
  fmap (Vec.take . Vec.concat . takeEnough @n . squeeze . absorb)
    . cut
    . uncurry (<>)
    . (id &&& pad r . fromIntegral . V.length)
  where
    cut :: Vector Bit -> Maybe [Vec ('S r) Bool]
    cut p =
      if V.null p
        then pure []
        else
          fmap (uncurry (:))
            . bitraverse
              (Vec.fromList . V.toList . V.map unBit)
              cut
            $ V.splitAt (Nat.reflectToNum r) p
    absorb :: [Vec ('S r) Bool] -> State l
    absorb =
      foldr
        ( \p_i s ->
            f (s `xor` restructureVec l (p_i Vec.++ pure @(Vec ('S c)) False))
        )
        . pure
        $ pure zeroBits
    squeeze :: State l -> Stream (Vec ('S r) Bool)
    squeeze =
      embed
        . ( Vec.take @('S r) . unstructureVec l
              &&& ana ((Vec.take @('S r) . unstructureVec l &&& id) . f)
          )

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- __FIXME__: See "Sponge" for the justification for these.
{-# OPTIONS_GHC -Wno-missing-safe-haskell-mode -Wno-unsafe #-}
-- Because of big `Fin`s.
{-# OPTIONS_GHC -freduction-depth=0 #-}
-- | This is an implementation of [Keccak](https://keccak.team/keccak.html) with
--   a lot of static guarantees.
--
--   There are some deviations from the spec:
-- * `keccak_f`, rather than taking the value of @b@ takes the value of @l@, so
--   where the spec says "Keccak-/f/[1600]", we would write
--   @`keccak_f` (Proxy \@(`Nat.fromGHC` 6))@.
-- * `keccak_rc` and `keccak_c` each take an extra `Proxy` (@l@ and @r@,
--   respectively), but these are uniquely determined by the other parameters,
--   so it's a minor annoyance.
module Keccak
  ( keccak_rc,
    keccak_c,
    keccak,
  )
where
import Control.Applicative (Applicative (..))
import Control.Category (Category (..))
import Control.Lens ((&), (.~))
import Data.Bit (Bit (..), Vector)
import Data.Bits (Bits (..))
import Data.Bool (Bool (..), bool)
import Data.Eq (Eq (..))
import Data.Fin (Fin)
import Data.Foldable (Foldable (..))
import Data.Function (flip)
import Data.Maybe (Maybe)
import Data.Ord (Ord (..))
import Data.Proxy (Proxy (..))
import Data.Semigroup (Semigroup (..))
import Data.Type.Nat (Nat (..))
import Data.Type.Nat qualified as Nat
import Data.Type.Nat.LE qualified as Nat
import Data.Vec.Lazy (Vec (..), (!))
import Data.Vec.Lazy qualified as Vec
import Data.Vec.Lazy.Lens qualified as Vec
import Data.Vector.Generic qualified as V
import Data.Word (Word64)
import Numeric (Floating (..))
import "this" Sponge (B, Block, State, W, WBits, sponge)
import Prelude
  ( Double,
    Integral (..),
    Num (..),
    RealFrac (truncate),
    fromIntegral,
    ($),
  )
imap2 :: (Fin m -> Fin n -> a -> b) -> Vec m (Vec n a) -> Vec m (Vec n b)
imap2 f = Vec.imap $ Vec.imap . f
-- = Keccak-specific definitions
theta :: Nat.SNatI (W l) => Proxy l -> State l -> State l
theta Proxy a =
  let c =
        Vec.map
          ( \ax ->
              Vec.tabulate (\z -> foldr1 xor (Vec.tabulate (\y -> (ax ! y) ! z)))
          )
          a
      d =
        Vec.tabulate
          ( \x ->
              Vec.tabulate
                ( \i ->
                    let q = (c ! (x - 1)) ! i
                        u = (c ! (x + 1)) ! (i - 1)
                     in q `xor` u
                )
          )
   in imap2 (\x _ -> xor (d ! x)) a
-- |
--  __FIXME__: I’m pretty sure this isn’t doing right thing.
rho_pi :: Nat.SNatI (W l) => Proxy l -> State l -> State l
rho_pi Proxy a =
  Vec.tabulate
    ( \y -> do
        Vec.ifoldr
          ( \x ax by ->
              by & Vec.ix (2 * x + 3 * y)
                .~ Vec.tabulate (\i -> (ax ! y) ! (i - rot_tbl x y))
          )
          (pure $ pure False)
          a
    )
chi :: Bits a => Vec Block a -> Vec Block a
chi a = Vec.imap (\x -> xor $ complement (a ! (x + 1)) .&. (a ! (x + 2))) a
iota :: Nat.SNatI (W l) => Proxy l -> Word64 -> State l -> State l
iota Proxy rc =
  Vec.imap
    ( \x ax ->
        bool
          ax
          ( Vec.imap
              ( \y axy ->
                  bool axy (Vec.imap (xor . get_round_bit rc) axy) $ y == 0
              )
              ax
          )
          $ x == 0
    )
round ::
  (Nat.SNatI (W l), Bits (WBits l)) => Proxy l -> Word64 -> State l -> State l
round l rc = iota l rc . chi . rho_pi l . theta l
-- Should calculate this rather than building an explicit `Vec`.
round_consts :: Nat.LE n_r (Nat.FromGHC 24) => Vec n_r Word64
round_consts =
  Vec.take
    ( 0x00000001
        ::: 0x00008082
        ::: 0x0000808a
        ::: 0x80008000
        ::: 0x0000808b
        ::: 0x80000001
        ::: 0x80008081
        ::: 0x00008009
        ::: 0x0000008a
        ::: 0x00000088
        ::: 0x80008009
        ::: 0x8000000a
        ::: 0x8000808b
        ::: 0x800000000000008b
        ::: 0x8000000000008089
        ::: 0x8000000000008003
        ::: 0x8000000000008002
        ::: 0x8000000000000080
        ::: 0x800000000000800a
        ::: 0x800000008000000a
        ::: 0x8000000080008081
        ::: 0x8000000080008080
        ::: 0x0000000080000001
        ::: 0x8000000080008008
        ::: VNil
    )
-- where
--   rc t = (x ** t `mod` x ** 8 + x ** 6 + x ** 5 + x ** 4 + 1) `mod` x
rot_tbl :: forall w. Nat.SNatI w => Fin Block -> Fin Block -> Fin w
rot_tbl x y =
  let m :: Vec Block (Vec Block (Fin w))
      m =
        (0 ::: 36 ::: 3 ::: 105 ::: 210 ::: VNil)
          ::: (1 ::: 300 ::: 10 ::: 45 ::: 66 ::: VNil)
          ::: (190 ::: 6 ::: 171 ::: 15 ::: 253 ::: VNil)
          ::: (28 ::: 55 ::: 153 ::: 21 ::: 120 ::: VNil)
          ::: (91 ::: 276 ::: 231 ::: 136 ::: 78 ::: VNil)
          ::: VNil
   in (m ! x) ! y
get_round_bit :: Nat.SNatI w => Word64 -> Fin w -> Bool
get_round_bit round_c bit_i =
  let the_bit =
        round_c .&. truncate (2 ** fromIntegral bit_i :: Double)
   in the_bit > 0
-- | Ideally, this function would accept @b@ rather than @l@ as a type
--   parameter, but it's hard to define the type that way, so this accepts @l@,
--   which is required to be in the range [0..6], producing the following
--   values:
--
--   +---+-----+----+------+
--   | l | n_r | w  |  b   |
--   +===+=====+====+======+
--   | 0 |  12 |  1 |   25 |
--   +---+-----+----+------+
--   | 1 |  14 |  2 |   50 |
--   +---+-----+----+------+
--   | 2 |  16 |  4 |  100 |
--   +---+-----+----+------+
--   | 3 |  18 |  8 |  200 |
--   +---+-----+----+------+
--   | 4 |  20 | 16 |  400 |
--   +---+-----+----+------+
--   | 5 |  22 | 32 |  800 |
--   +---+-----+----+------+
--   | 6 |  24 | 64 | 1600 |
--   +---+-----+----+------+
keccak_f ::
  forall l n_r.
  ( Nat.SNatI (W l),
    Bits (WBits l),
    n_r ~ Nat.Plus (Nat.FromGHC 12) (Nat.Mult2 l),
    -- This implies @l <= 6@
    Nat.LE n_r (Nat.FromGHC 24)
  ) =>
  Proxy l ->
  State l ->
  State l
keccak_f l a = foldl (flip $ round l) a $ round_consts @n_r
pad10x1 :: Nat.SNatI x => Proxy ('S x) -> Nat -> Vector Bit
pad10x1 x l =
  -- __FIXME__: This is flagged as non-exhaustive, but because the `Proxy` is
  --            explicitly `'S`, it _is_ exhaustive.
  let (S x') = Nat.reflect x
   in V.cons (Bit True) $
        V.replicate
          ( fromIntegral $
              case x' - (l `mod` S x') of
                Z -> x'
                S neededZeros -> neededZeros
          )
          (Bit False)
          <> V.singleton (Bit True)
-- | Ideally this wouldn't require @l@, but until `keccak_f` accepts @b@ instead
--   of @l@, this is necessary.
keccak_rc ::
  forall l' l n_r r c n.
  ( Nat.SNatI r,
    Nat.SNatI c,
    Nat.SNatI n,
    Nat.SNatI (W l),
    Bits (WBits l),
    Nat.SNatI (Nat.Mult Block (W l)),
    B l ~ Nat.Plus ('S r) ('S c),
    Nat.LE ('S r) (B l),
    Nat.LE ('S l') (Nat.Mult ('S n) ('S r)),
    n_r ~ Nat.Plus (Nat.FromGHC 12) (Nat.Mult2 l),
    -- This implies @l <= 6@
    Nat.LE n_r (Nat.FromGHC 24)
  ) =>
  Proxy l ->
  Proxy ('S r) ->
  Proxy ('S c) ->
  -- | This must be at least ⌈l' / r⌉
  Proxy ('S n) ->
  Vector Bit ->
  Maybe (Vec ('S l') Bool)
keccak_rc l = sponge l (keccak_f l) pad10x1
-- | This taking @r@ seems a bit silly, but since @b@ is fixed, it means that
--   @r@ is uniquely determined by @c@ (and vice versa).
keccak_c ::
  forall l' l r c n.
  ( l ~ Nat.FromGHC 6,
    Nat.SNatI r,
    Nat.SNatI c,
    Nat.SNatI n,
    B l ~ Nat.Plus ('S r) ('S c),
    Nat.LE ('S r) (B l),
    Nat.LE ('S l') (Nat.Mult ('S n) ('S r))
  ) =>
  Proxy ('S r) ->
  Proxy ('S c) ->
  -- | This must be at least ⌈l' / r⌉
  Proxy ('S n) ->
  Vector Bit ->
  Maybe (Vec ('S l') Bool)
keccak_c = keccak_rc (Proxy @l)
keccak ::
  forall l' l r c n.
  ( l ~ Nat.FromGHC 6,
    r ~ Nat.FromGHC 1024,
    c ~ Nat.FromGHC 576,
    Nat.SNatI n,
    Nat.LE r (B l),
    Nat.LE ('S l') (Nat.Mult ('S n) r)
  ) =>
  -- | This must be at least ⌈l' / 1024⌉
  Proxy ('S n) ->
  Vector Bit ->
  Maybe (Vec ('S l') Bool)
keccak = keccak_c (Proxy @r) (Proxy @c)

  ./keccak/keccak.cabal