1	import Base
2	import ../libs/AES256GCMCore.bend as Core
3	
4	# Compose short, independently checked computation steps. Concrete AES vectors
5	# can reuse these equalities without normalizing every round in one kernel call.
6	
7	def next_word(+i: U32, rcon: U32, +words: List<&2, U32>) -> U32:
8	    previous = Core.aes_word_at(words, U32.to_nat(U32.sub(i, 1)))
9	    back8 = Core.aes_word_at(words, U32.to_nat(U32.sub(i, 8)))
10	    part = Core.aes_key_word_kind(previous, rcon, U32.is_eq(U32.mod(i, 8), 0), U32.is_eq(U32.mod(i, 8), 4))
11	    U32.xor(back8, part)
12	
13	def expand_step(+fuel: Nat, +i: U32, +rcon: U32, +words: List<&2, U32>,
14	    +next_words: List<&2, U32>, +next_rcon: U32,
15	    words_match: {List.append(&2, U32, words, [next_word(i, rcon, words)]) == next_words : List<&2, U32>},
16	    rcon_match: {Core.aes_rcon_kind(rcon, U32.is_eq(U32.mod(i, 8), 0)) == next_rcon : U32}) ->
17	    {Core.aes_expand.go(1n+fuel, i, rcon, words) ==
18	     Core.aes_expand.go(fuel, U32.add(i, 1), next_rcon, next_words) : List<&2, U32>}:
19	    +actual_rcon = Core.aes_rcon_kind(rcon, U32.is_eq(U32.mod(i, 8), 0))
20	    Equal.trans(List<&2, U32>, Core.aes_expand.go(1n+fuel, i, rcon, words),
21	        Core.aes_expand.go(fuel, U32.add(i, 1), actual_rcon, next_words),
22	        Core.aes_expand.go(fuel, U32.add(i, 1), next_rcon, next_words),
23	        Equal.cong(List<&2, U32>, List<&2, U32>,
24	            value => Core.aes_expand.go(fuel, U32.add(i, 1), actual_rcon, value),
25	            List.append(&2, U32, words, [next_word(i, rcon, words)]), next_words, words_match),
26	        Equal.cong(U32, List<&2, U32>, value => Core.aes_expand.go(fuel, U32.add(i, 1), value, next_words),
27	            actual_rcon, next_rcon, rcon_match))
28	
29	def rounds_step(+fuel: Nat, +state: List<&2, U32>, +words: List<&2, U32>,
30	    +round: Nat, +next_state: List<&2, U32>,
31	    same: {Core.aes_middle_round(state, words, round) == next_state : List<&2, U32>}) ->
32	    {Core.aes_rounds.go(1n+fuel, state, words, round) ==
33	     Core.aes_rounds.go(fuel, next_state, words, (round + 1n : Nat)) : List<&2, U32>}:
34	    Equal.cong(List<&2, U32>, List<&2, U32>,
35	        value => Core.aes_rounds.go(fuel, value, words, (round + 1n : Nat)),
36	        Core.aes_middle_round(state, words, round), next_state, same)
37	
38	def expand_step_at(+fuel: Nat, +i: U32, +rcon: U32, +words: List<&2, U32>,
39	    +next_words: List<&2, U32>, +next_rcon: U32, +next_i: U32,
40	    words_match: {List.append(&2, U32, words, [next_word(i, rcon, words)]) == next_words : List<&2, U32>},
41	    rcon_match: {Core.aes_rcon_kind(rcon, U32.is_eq(U32.mod(i, 8), 0)) == next_rcon : U32},
42	    index_match: {U32.add(i, 1) == next_i : U32}) ->
43	    {Core.aes_expand.go(1n+fuel, i, rcon, words) ==
44	     Core.aes_expand.go(fuel, next_i, next_rcon, next_words) : List<&2, U32>}:
45	    Equal.trans(List<&2, U32>, Core.aes_expand.go(1n+fuel, i, rcon, words),
46	        Core.aes_expand.go(fuel, U32.add(i, 1), next_rcon, next_words),
47	        Core.aes_expand.go(fuel, next_i, next_rcon, next_words),
48	        expand_step(fuel, i, rcon, words, next_words, next_rcon, words_match, rcon_match),
49	        Equal.cong(U32, List<&2, U32>, value => Core.aes_expand.go(fuel, value, next_rcon, next_words),
50	            U32.add(i, 1), next_i, index_match))
51	
52	def expand_step_number_at(+fuel: Nat, +remaining: Nat,
53	    fuel_match: {remaining == 1n+fuel : Nat},
54	    +i: U32, +rcon: U32, +words: List<&2, U32>,
55	    +next_words: List<&2, U32>, +next_rcon: U32, +next_i: U32,
56	    words_match: {List.append(&2, U32, words, [next_word(i, rcon, words)]) == next_words : List<&2, U32>},
57	    rcon_match: {Core.aes_rcon_kind(rcon, U32.is_eq(U32.mod(i, 8), 0)) == next_rcon : U32},
58	    index_match: {U32.add(i, 1) == next_i : U32}) ->
59	    {Core.aes_expand.go(remaining, i, rcon, words) ==
60	     Core.aes_expand.go(fuel, next_i, next_rcon, next_words) : List<&2, U32>}:
61	    Equal.trans(List<&2, U32>, Core.aes_expand.go(remaining, i, rcon, words),
62	        Core.aes_expand.go(1n+fuel, i, rcon, words),
63	        Core.aes_expand.go(fuel, next_i, next_rcon, next_words),
64	        Equal.cong(Nat, List<&2, U32>, steps => Core.aes_expand.go(steps, i, rcon, words),
65	            remaining, 1n+fuel, fuel_match),
66	        expand_step_at(fuel, i, rcon, words, next_words, next_rcon, next_i,
67	            words_match, rcon_match, index_match))
68	
69	def rounds_fuel_matches(+left: Nat, +right: Nat, +state: List<&2, U32>,
70	    +words: List<&2, U32>, +round: Nat, same: {left == right : Nat}) ->
71	    {Core.aes_rounds.go(left, state, words, round) ==
72	     Core.aes_rounds.go(right, state, words, round) : List<&2, U32>}:
73	    Equal.cong(Nat, List<&2, U32>, fuel => Core.aes_rounds.go(fuel, state, words, round), left, right, same)
74	
75	def rounds_step_at(+fuel: Nat, +state: List<&2, U32>, +words: List<&2, U32>,
76	    +round: Nat, +next_state: List<&2, U32>, +next_round: Nat,
77	    same: {Core.aes_middle_round(state, words, round) == next_state : List<&2, U32>},
78	    round_match: {(round + 1n : Nat) == next_round : Nat}) ->
79	    {Core.aes_rounds.go(1n+fuel, state, words, round) ==
80	     Core.aes_rounds.go(fuel, next_state, words, next_round) : List<&2, U32>}:
81	    Equal.trans(List<&2, U32>, Core.aes_rounds.go(1n+fuel, state, words, round),
82	        Core.aes_rounds.go(fuel, next_state, words, (round + 1n : Nat)),
83	        Core.aes_rounds.go(fuel, next_state, words, next_round),
84	        rounds_step(fuel, state, words, round, next_state, same),
85	        Equal.cong(Nat, List<&2, U32>,
86	            value => Core.aes_rounds.go(fuel, next_state, words, value),
87	            (round + 1n : Nat), next_round, round_match))