{-# LANGUAGE ScopedTypeVariables, LambdaCase #-}
{-
Generate the Keccak round constants from its formula.

Copyright 2018 Ken Takusagawa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or (at
your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

-}

module Main where {
import System.Environment(getArgs);
import Control.Exception(assert);
import Debug.Trace(trace);
import Data.Function((&));
import Control.Category((>>>));
import Prelude hiding((.),(>>));
import Data.Bits(xor);
--import System.IO(hSetBuffering,stdout,BufferMode(LineBuffering));
import Data.List;
--import Control.Monad;
--import Data.Maybe;
import qualified Data.Map as Map;
import Data.Map(Map);
--import qualified Data.Set as Set;
--import Data.Set(Set);
import qualified Data.Bifunctor as Bifunctor;

-- to avoid the redundancy warning
trace_placeholder :: ();
trace_placeholder = (trace,assert) & (id >>> error) "trace_placeholder";

main :: IO();
main = getArgs >>= \case{
["bin"] -> round_constants_in_binary 0;
["bin",num_extra_rounds] -> round_constants_in_binary (read num_extra_rounds);
["hex"] -> round_constants_in_hex ( secure_num_rounds standard_l);
["hex",total_rounds] -> round_constants_in_hex (read total_rounds);
["lfsr"] -> rc_chunks & unwords & putStrLn;
_ -> undefined;
};

lfsr :: [Bool] -> [Bool];
lfsr old = let {
r :: [Bool];
r = False : old; -- left shift
r8 :: Bool;
r8 = r!!8; -- msb that just got shifted "off"
alter :: Bool -> Bool;
alter = xor r8;
} in
zipWith ($) [alter, id, id, id, alter, alter, alter, id] r;  -- this does the truncation to 8 bits because zipWith takes the shorter

initial :: [Bool];
initial = (True : repeat False) & take 8;

-- starting with zero repeats with period 1 (i.e., does not change)
-- anything else repeats with cycle 255
as_char :: Bool -> Char;
as_char True = '1';
as_char False = '.';

as_string :: [Bool] -> String;
as_string = map as_char >>> nchunks 4 >>> intersperse " " >>> concat;

-- the lowercase rc function.
rc_get_bit_as_specified :: Integer -> Bool;
rc_get_bit_as_specified t = case mod t 255 of {
0 -> True;
tmod -> initial & iterate lfsr & (\l -> genericIndex l tmod) & head
};

-- because the LFSR naturally cycles with a period of 255, there is no need to explicitly mod 255 (actually, we do need to handle negative numbers)
-- and because the LSB of initial is 1, there is no need to special case 0
rc_get_bit :: Integer -> Bool;
rc_get_bit t = initial & iterate lfsr & (\l -> genericIndex l (mod t 255)) & head;

-- the 7 is because 64-bits, there are 7 possible positions a bit can be set
set_bits :: Integer -> Integer -> (Integer,Bool); -- bit position and value;
set_bits round_index j = (2^j-1, rc_get_bit (j+7*round_index));

-- 24 rounds, only 168 bits needed total

do_set_bit :: [a] -> (Integer, a) -> [a];
do_set_bit old (position,new) = let {
(left,right) = genericSplitAt position old;
} in left ++ (new : tail right);

zero :: [Bool];
zero = repeat False;

round_constant_infinite_list :: Integer -> Integer -> [Bool];
round_constant_infinite_list l round_index = enumFromTo 0 l & map (set_bits round_index) & foldl' do_set_bit zero;

standard_l :: Integer;
standard_l = 6;

round_constant_as_specified :: Integer -> [Bool];
round_constant_as_specified round_index = let {
l = standard_l -- hardcode 64-bit max
} in round_constant_infinite_list l round_index & take (2^l);

to_integer :: [Bool] -> Integer;
to_integer = reverse >>> foldl' (\accum b -> 2*accum+if b then 1 else 0) 0;

to_hex :: Integer -> Char;
to_hex i = fromInteger i + (if i<10 then fromEnum '0' else (fromEnum 'a' - 10)) & toEnum;

to_hex_string :: [Bool] -> String;
to_hex_string = nchunks 4 >>> map to_integer >>> map to_hex >>> reverse;

to_64bit_hex :: String -> String;
to_64bit_hex = reverse >>> (\x -> x ++ (repeat '0')) >>> take 16 >>> reverse >>> (\x -> "0x"++x);

put :: [Bool] -> IO ();
put = as_string >>> putStrLn;

standard_round_constants :: Integer -> [(Integer,[Bool])];
standard_round_constants total_rounds = let {
 indices = round_indices standard_l & genericTake total_rounds & reverse;
} in map (\i -> (i, round_constant_as_specified i)) indices;

round_constants_in_binary :: Integer -> IO ();
round_constants_in_binary extra = standard_round_constants (extra + secure_num_rounds standard_l)
& mapM_ (Bifunctor.second (map as_char) >>> uncurry numbered_line);

numbered_line :: Integer -> String -> IO ();
numbered_line n s = putStrLn $ show n ++ " " ++ s;

round_constants_in_hex :: Integer -> IO ();
round_constants_in_hex total_rounds = standard_round_constants total_rounds
& mapM_ (Bifunctor.second (to_hex_string >>> ("0x"++)) >>> uncurry numbered_line);

nchunks :: Integer -> [a] -> [[a]];
nchunks n = unfoldr (\l -> if null l then Nothing else Just $ genericSplitAt n l);

-- read off the LFSR 7 bits at a time
rc_chunks :: [String];
rc_chunks = map rc_get_bit [0..] & take 255 & map as_char & nchunks 7;

-- the round constants repeat with a cycle of 255, only because gcd 255 and 7 = 12
-- if chunks of 5 then period div 255 5 = 51 etc.

make_map_with_repeats :: Ord k => [(k,v)] -> Map k [v];
make_map_with_repeats = map (\(k,v) -> (k,[v])) >>> Map.fromListWith (++);

has_repeats :: (a,[b]) -> Bool;
has_repeats (_,l) = length l > 1;

-- rounds 5 and 22 are repeats, and 6 and 20
-- [-3,3],[-1,8] for 3 and 1 extra rounds

-- if 32-bit, 22 rounds
-- [[6,20],[11,19]]
-- [[-6,6,20],[-5,15],[-3,3],[-2,16],[-1,8],[11,19]]

-- 16 bit, 20 rounds
-- [[0,5],[4,12],[7,10],[11,19]]
-- [[-6,-4,6],[-5,15],[-3,3],[-2,16],[-1,8],[0,5],[4,12],[7,10],[11,19]]

-- 8 bit, 18 rounds
-- [[0,5],[2,8],[4,12,13],[7,10]]

-- controversial formula that got changed to 12+l then changed back
secure_num_rounds :: Integer -> Integer;
secure_num_rounds l = 12+2*l;

-- the round indices go negative, or start not at 0
round_indices :: Integer -> [Integer];
round_indices = secure_num_rounds >>> pred >>> iterate pred;

repeats_in_round_constants :: Integer -> Integer -> [[Integer]];
repeats_in_round_constants l num_rounds_extra = let {
 num_rounds = secure_num_rounds l + num_rounds_extra;
 indices = round_indices l;
} in indices
& genericTake num_rounds
& map ( round_constant_infinite_list l >>> take (2^l))
& \x -> zip x indices
& make_map_with_repeats
& Map.toList
& filter has_repeats
& map (snd >>> sort)
& sort;

} --end