Theory MLWE

theory MLWE

imports Lemmas_for_spmf
        "Game_Based_Crypto.CryptHOL_Tutorial"

begin

section ‹Module Learning-with-Errors Problem (module-LWE)›

text ‹‹Berlekamp_Zassenhaus› loads the vector type ‹'a vec› from ‹Jordan_Normal_Form.Matrix›. 
This doubles the symbols ‹\$› and ‹χ› for ‹vec_nth› and ‹vec_lambda›. Thus we delete the
‹vec_index› for type ‹'a vec›. Still some type ambiguities remain.›

unbundle %invisible lifting_syntax
no_adhoc_overloading %invisible Monad_Syntax.bind bind_pmf
declare %invisible [[names_short]]


text ‹Here the actual theory starts.›

text ‹We introduce a locale ‹module_lwe› that represents the module-Learning-with-Errors 
(module-LWE) problem in the setting of Kyber. 
The locale takes as input:
\begin{itemize}
\item ‹type_a› the type of the quotient ring of Kyber. (This is a side effect of the Harrison 
    trick in the Kyber locale.)
\item ‹type_k› the finite type for indexing vectors in Kyber. The cardinality is exactely $k$.
   (This is a side effect of the Harrison trick in the Kyber locale.)
\item ‹idx› an indexing function from ‹'k› to ‹{0..<k}›
\item ‹eta› the specification value for the centered binomial distribution $\beta_\eta$
\end{itemize}›

locale module_lwe = 
fixes type_a :: "('a :: qr_spec) itself" 
  and type_k :: "('k ::finite) itself"
  and k :: nat
  and idx :: "'k::finite ⇒ nat"
  and eta :: nat
assumes "k = CARD('k)"
  and bij_idx: "bij_betw idx (UNIV::'k set) {0..<k}"
 
begin

text ‹The adversary in the module-LWE takes a matrix ‹A::(('b, 'n) vec, 'm) vec› and a vector
  ‹t::('b, 'm) vec› and returns a probability distribution on ‹bool› 
  guessing whether the given input was randomly generated or a valid module-LWE instance.›
type_synonym ('b, 'n, 'm) adversary = 
  "(('b, 'n) vec, 'm) vec ⇒ ('b, 'm) vec ⇒  bool spmf"

text ‹Next, we want to define the centered binomial distributions $\beta_\eta$.
‹bit_set› returns the set of all bit lists of length ‹eta›.
‹beta› is the centered binomial distribution $\beta_\eta$ as a ‹pmf› on the quotient ring $R_q$.
‹beta_vec› is then centered binomial distribution $\beta_\eta ^k$ on vectors in $R_q^k$.›

definition bit_set :: "int list set" where
"bit_set = {xs:: int list. set xs ⊆ {0,1} ∧ length xs = eta}"

lemma finite_bit_set:
"finite bit_set"
unfolding bit_set_def 
by (simp add: finite_lists_length_eq)

lemma bit_set_nonempty:
"bit_set ≠ {}"
proof -
  have "replicate eta 0 ∈ bit_set" unfolding bit_set_def by auto
  then show ?thesis by auto
qed

adhoc_overloading %invisible Monad_Syntax.bind bind_pmf 
definition beta :: "'a qr pmf" where
"beta = do {
    as ← pmf_of_set (bit_set);
    bs ← pmf_of_set (bit_set);
    return_pmf (to_module (∑i<eta. as ! i - bs! i))
  } "

definition beta_vec :: "('a qr , 'k) vec pmf" where
"beta_vec = do {
    (xs :: 'a qr list) ← replicate_pmf (k) (beta);
    return_pmf (χ i.  xs ! (idx i))
  }"

text ‹Since we work over ‹spmf›, we need to show that ‹beta_vec› is lossless.›

lemma lossless_beta_vec[simp]:
  "lossless_spmf (spmf_of_pmf beta_vec)"
by (simp)

text ‹We define the game versions of module-LWE.
Given an adversary ‹𝒜›, we have two games: 
in ‹game›, the instance given to the adversary is a module-LWE instance, 
whereas in ‹game_random›, the instance is chosen randomly.›

definition game :: "('a qr,'k,'k) adversary ⇒ bool spmf" where
  "game 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 'k) vec, 'k) vec set);
    s ← beta_vec;
    e ← beta_vec;
    b' ← 𝒜 A (A *v s + e);
    return_spmf (b')
  }"

definition game_random :: "('a qr,'k,'k) adversary ⇒ bool spmf" where
  "game_random 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 'k) vec, 'k) vec set);
    b ← spmf_of_set (UNIV:: ('a qr, 'k) vec set);
    b' ← 𝒜 A b;
    return_spmf (b')
  }"

text ‹The advantage of an adversary ‹𝒜› returns a value how good the adversary is at guessing 
whether the instance is generated by the module-LWE or uniformly at random.›

definition advantage :: "('a qr,'k,'k) adversary ⇒ real" where 
"advantage 𝒜 = ¦spmf (game 𝒜) True - spmf (game_random 𝒜) True ¦"


text ‹Since the reduction proof of Kyber uses the module-LWE problem for two different dimensions
(ie.\ $A\in R_q^{(k+1)\times k}$ and $A\in R_q^{k\times k}$), we need a second definition of
the index function, the centered binomial distribution, the game and random game, and the advantage.
Here the problem is that the dimension $k$ of the vectors is hard-coded in the type ‹'k›. 
That makes it hard to ``simply add'' another dimension. A trick how this can be formalised 
nevertheless is to use the option type on ‹'k› to encode a type with $k+1$ elements.
With the option type, we can embed a vector of dimension $k$ indexed by the type ‹'k› into
a vector of dimension $k+1$ by adding a value for the index ‹None› (an element ‹a :: 'k› 
is mapped to ‹Some a›).
Note also that the additional index appears only in one dimension of $A$, resulting in 
a non-quadratic matrix.›

text ‹Index function of the option type ‹'k option›.›

fun idx' :: "'k option ⇒ nat" where
  "idx' None = 0" |
  "idx' (Some x) = idx x + 1"

lemma idx': "((x # xs) ! idx' i) = 
  ( if i = None then x else xs ! idx (the i))" 
  if "length xs = k" for xs and i::"'k option"
proof (cases i)
  case None
  then show ?thesis using nth_append_length[of xs x "[]"] that by (simp add: if_distrib)
next
  case (Some a)
  have "idx a < k" using bij_idx 
  by (meson UNIV_I atLeastLessThan_iff bij_betwE) 
  then show ?thesis using Some that by (simp add: if_distrib nth_append)
qed


lemma idx'_lambda:
  "(χ i. (x # xs) ! idx' i) = 
   (χ i. if i = None then x else xs ! idx (the i))" 
   if "length xs = k" for xs using idx'[OF that]
 by presburger

text ‹Definition of the centered binomial distribution $\beta_\eta^{k+1}$ and lossless property.›

definition beta_vec' :: "('a qr , 'k option) vec spmf" where
"beta_vec' = do {
    (xs :: 'a qr list) ← replicate_spmf (k+1) (beta);
    return_spmf (χ i.  xs ! (idx' i))
  }"

lemma lossless_beta_vec'[simp]:
  "lossless_spmf beta_vec'"
unfolding beta_vec'_def beta_def
by (simp add: replicate_spmf_def)

text ‹Some lemmas on replicate.›

lemma replicate_pmf_same_length:
assumes "⋀ xs. length xs = m ⟹ f xs = g xs"
shows "bind_pmf (replicate_pmf m p) f = bind_pmf (replicate_pmf m p) g"
by (metis (mono_tags, lifting) assms bind_pmf_cong mem_Collect_eq set_replicate_pmf)

lemma replicate_spmf_same_length:
assumes "⋀ xs. length xs = m ⟹ f xs = g xs"
shows "(replicate_spmf m p ⤜ f) = (replicate_spmf m p ⤜ g)"
unfolding replicate_spmf_def bind_spmf_of_pmf 
by (simp add: replicate_pmf_same_length[OF assms])

text ‹Lemma to split the ‹replicate (k+1)› function in ‹beta_vec'› into two parts: 
‹replicate k› and a separate value. Note, that the ‹xs› in the ‹do› notation below are 
always of length $k$. ›

no_adhoc_overloading Monad_Syntax.bind bind_pmf

lemma beta_vec':
  "beta_vec' = do {
    (xs :: 'a qr list) ← replicate_spmf (k) (beta);
    (x :: 'a qr) ← spmf_of_pmf beta;
    return_spmf (χ i.  if i = None then x else xs ! (idx (the i)))
  }"                                                    
unfolding beta_vec'_def idx'_lambda[symmetric] replicate_spmf_Suc_cons 
  bind_spmf_assoc return_bind_spmf
by (subst replicate_spmf_same_length[where 
  f = "(λy. spmf_of_pmf beta ⤜ (λya. return_spmf (vec_lambda (λi. (ya # y) ! idx' i))))" and
  g = "(λxs. spmf_of_pmf beta ⤜ (λx. return_spmf (vec_lambda (λi. if i = None then x else 
        xs ! idx (the i)))))"])
   (simp_all add: idx'_lambda)

adhoc_overloading Monad_Syntax.bind bind_pmf

text ‹Definition of the two games for the option type.›

definition game' :: "('a qr,'k,'k option) adversary ⇒ bool spmf" where
  "game' 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 'k) vec, 'k option) vec set);
    s ← beta_vec;
    e ← beta_vec';
    b' ← 𝒜 A (A *v s + e);
    return_spmf (b')
  }"

definition game_random' :: "('a qr,'k,'k option) adversary ⇒ bool spmf" where
  "game_random' 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 'k) vec, 'k option) vec set);
    b ← spmf_of_set (UNIV:: ('a qr, 'k option) vec set);
    b' ← 𝒜 A b;
    return_spmf (b')
  }"

text ‹Definition of the advantage for the option type.›

definition advantage' :: "('a qr,'k,'k option) adversary ⇒ real" where 
"advantage' 𝒜 = ¦spmf (game' 𝒜) True - spmf (game_random' 𝒜) True ¦"

text ‹Game and random game for finite type with one element only›

definition beta1 :: "('a qr , 1) vec pmf" where
"beta1 = bind_pmf beta (λx. return_pmf (χ i. x))"

definition game1 :: "('a qr, 1, 1) adversary ⇒ bool spmf" where
  "game1 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 1) vec, 1) vec set);
    s ← spmf_of_pmf beta1;
    e ← spmf_of_pmf beta1;
    b' ← 𝒜 A (A *v s + e);
    return_spmf (b')
  }"

definition game_random1 :: "('a qr,1,1) adversary ⇒ bool spmf" where
  "game_random1 𝒜  = do {
    A ← spmf_of_set (UNIV:: (('a qr, 1) vec, 1) vec set);
    b ← spmf_of_set (UNIV:: ('a qr, 1) vec set);
    b' ← 𝒜 A b;
    return_spmf (b')
  }"

text ‹The advantage of an adversary ‹𝒜› returns a value how good the adversary is at guessing 
whether the instance is generated by the module-LWE or uniformly at random.›

definition advantage1 :: "('a qr,1,1) adversary ⇒ real" where 
"advantage1 𝒜 = ¦spmf (game1 𝒜) True - spmf (game_random1 𝒜) True ¦"

end
end