/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Prolog implementation of ChaCha20.
   Written Dec. 2021 by Markus Triska (triska@metalevel.at)

   For more information, see "ChaCha, a variant of Salsa20" by
   Daniel J. Bernstein: https://cr.yp.to/chacha/chacha-20080128.pdf

   Tested with Scryer Prolog.
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

:- use_module(library(format)).
:- use_module(library(dcgs)).
:- use_module(library(lists)).
:- use_module(library(between)).
:- use_module(library(lambda)).
:- use_module(library(clpz)).

/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

?- numlist(0, 31, Ks),
   chacha20_init(Ks, [0,0,0,9,0,0,0,74,0,0,0,0], [1,0,0,0], As0),
   chacha20(As0, As),
   format("~n~nResult:~n~n", []),
   hex_array(As).
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   ChaCha updates a, b, c, d as follows:

   a += b; d ^= a; d <<<= 16;
   c += d; b ^= c; b <<<= 12;
   a += b; d ^= a; d <<<= 8;
   c += d; b ^= c; b <<<= 7;
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

quarter_round(Vals0, Vals) :-
        Is = [a = a+b, d = d xor a, d = rotl(d,16),
              c = c+d, b = b xor c, b = rotl(b,12),
              a = a+b, d = d xor a, d = rotl(d,8),
              c = c+d, b = b xor c, b = rotl(b,7)],
        phrase(interpret_instructions(Is), [Vals0], [Vals]).

quarter_rounds([ % first the columns:
                [a,=,=,=,
                 b,=,=,=,
                 c,=,=,=,
                 d,=,=,=],
                [=,a,=,=,
                 =,b,=,=,
                 =,c,=,=,
                 =,d,=,=],
                [=,=,a,=,
                 =,=,b,=,
                 =,=,c,=,
                 =,=,d,=],
                [=,=,=,a,
                 =,=,=,b,
                 =,=,=,c,
                 =,=,=,d],
               % then the diagonals
                [a,=,=,=,
                 =,b,=,=,
                 =,=,c,=,
                 =,=,=,d],
                [=,a,=,=,
                 =,=,b,=,
                 =,=,=,c,
                 d,=,=,=],
                [=,=,a,=,
                 =,=,=,b,
                 c,=,=,=,
                 =,d,=,=],
                [=,=,=,a,
                 b,=,=,=,
                 =,c,=,=,
                 =,=,d,=]
                ]).


interpret_instructions([]) --> [].
interpret_instructions([I|Is]) -->
        interpret(I),
        interpret_instructions(Is).

state0_state(S0, S), [S] --> [S0].

interpret(Var = Expr) -->
        expr_value(Expr, Value),
        state0_state(Vals, [Var=Value|Vals]).

expr_value(X0+Y0, V) -->
        var_value(X0, X),
        var_value(Y0, Y),
        { plus_(X, Y, V) }.
expr_value(X0 xor Y0, V) -->
        var_value(X0, X),
        var_value(Y0, Y),
        { maplist(xor_, X, Y, V) }.
expr_value(rotl(X0,R), V) -->
        var_value(X0, X),
        { rotl_(X, R, V) }.


var_value(Which, Value) -->
        state0_state(State, State),
        { once(member(Which=Value, State)) }.

xor_('0', '0', '0').
xor_('0', '1', '1').
xor_('1', '0', '1').
xor_('1', '1', '0').

/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
?- integer_word(5, Ws) ; false.
%@    Ws = "0000000000000000000 ..."
%@    Ws = "00000000000000000000000000000101"
%@ ;  false.

?- integer_word(5, Ws), rotl_(Ws, 1, Rs) ; false.
%@    Ws = "0000000000000000000 ...", Rs = "0000000000000000000 ..."
%@    Ws = "00000000000000000000000000000101", Rs = "00000000000000000000000000001010"
%@ ;  false.

- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

rotl_(Bs0, R, Bs) :-
        length(Ls, R),
        append(Ls, Rs, Bs0),
        append(Rs, Ls, Bs).


plus_(As, Bs, Ss) :-
        binary_decimal(As, DA),
        binary_decimal(Bs, BA),
        S #= (DA + BA) mod 2^32,
        integer_word(S, Ss).

chacha20(As0, As) :-
        format("starting: ~n~n", []),
        hex_array(As0),
        numlist(1, 10, Rounds),
        foldl(apply_round, Rounds, As0, As1),
        maplist(plus_, As0, As1, As).

apply_round(R, As0, As) :-
        format("~n~nRound ~d-~d~n~`=t~20+~n", [R*2-1,R*2]),
        quarter_rounds(Qs),
        foldl(apply_quarter_round, Qs, As0, As),
        hex_array(As).

apply_quarter_round(Q, As0, As) :-
        phrase(collect_values(Q, As0, As), Values0),
        quarter_round(Values0, Values),
        assign_results("abcd", Values).

assign_results([], _).
assign_results([Var|Vars], Values) :-
        once(member(Var=Value, Values)),
        once(member(result(Var,Value), Values)),
        assign_results(Vars, Values).

collect_values([], [], []) --> [].
collect_values([E|Es], [A0|As0], [A|As]) -->
        collect_(E, A0, A),
        collect_values(Es, As0, As).

collect_(=, X, X) --> [].
collect_(a, A, R) --> value_and_result(a, A, R).
collect_(b, B, R) --> value_and_result(b, B, R).
collect_(c, C, R) --> value_and_result(c, C, R).
collect_(d, D, R) --> value_and_result(d, D, R).

value_and_result(Var, Val, R) --> [Var=Val,result(Var,R)].

hex_array(As) :-
        phrase(hex_array(As), Ls),
        format("~s", [Ls]).

hex_array(As) -->
        { maplist(word_hex, As, Hs) },
        hex_lines(Hs).

hex_lines([]) --> [].
hex_lines([A,B,C,D|Ls]) -->
        format_("~s   ~s   ~s   ~s~n", [A,B,C,D]),
        hex_lines(Ls).

chacha20_init(Key, Nonce, Counter, As) :-
        As = [A,B,C,D,
              E,F,G,H,
              I,J,K,L,
              M,N,O,P],
        integer_word(0x61707865, A),
        integer_word(0x3320646e, B),
        integer_word(0x79622d32, C),
        integer_word(0x6b206574, D),
        bytes_words(Key, [E,F,G,H,I,J,K,L]),
        bytes_words(Counter, [M]),
        bytes_words(Nonce, [N,O,P]).


integer_word(C, W) :-
        phrase(format_("~`0t~2r~32+", [C]), W).

bytes_words([], []).
bytes_words([A,B,C,D|Bs], [W|Ws]) :-
        phrase(bytes([D,C,B,A]), W),
        bytes_words(Bs, Ws).

bytes([]) --> [].
bytes([B|Bs]) -->
        format_("~`0t~2r~8+", [B]),
        bytes(Bs).

binary_decimal(Bs, D) :-
        maplist(b_i, Bs, Is),
        foldl(plus_times2, Is, 0, D).

plus_times2(I, S0, S) :- S #= 2*S0 + I.

b_i('0', 0).
b_i('1', 1).

word_hex(Bytes, Hex) :-
        binary_decimal(Bytes, D),
        phrase(("0x",format_("~`0t~16r~8+", [D])), Hex).

words_bytes([]) --> [].
words_bytes([W|Ws]) -->
        word_bytes(W),
        words_bytes(Ws).

word_bytes([]) --> [].
word_bytes(Ws) -->
        { length(Bss, 4),
          maplist(\Ls^length(Ls, 8), Bss),
          phrase(seqq(Bss), Ws),
          maplist(binary_decimal, Bss, [A,B,C,D]) },
        [D,C,B,A].