let rewrite_instructions (prg: Types.p_exp) : Types.p_exp =
  (* function takes a program and replaces all occurrences of powermod with
     simpler instructions *)

  let Main (i, o, prg) = prg in

  let rec contains_rewrite (prg: Types.c_exp) : Types.a_exp option =
    (* if the ast contains powermod anywhere returns the powermod, otherwise
       returns none *)
    match prg with
    | Skip -> None
    | Assignment (_, a) -> contains_rewrite_a a
    | Sequence (c1, c2) -> (
        match contains_rewrite c1, contains_rewrite c2 with
        | None, None -> None
        | Some a, _
        | _, Some a -> Some a
      )
    | If (b, c1, c2) -> (
        match contains_rewrite_b b,
              contains_rewrite c1,
              contains_rewrite c2
        with
        | None, None, None -> None
        | Some a, _, _
        | _, Some a, _
        | _, _, Some a -> Some a
      )
    | While (b, c) -> (
        match contains_rewrite_b b, contains_rewrite c with
        | None, None -> None
        | Some a, _
        | _, Some a -> Some a
      )
    | For (c1, b, c2, c3) -> (
        match contains_rewrite c1,
              contains_rewrite_b b,
              contains_rewrite c2,
              contains_rewrite c3
        with
        | None, None, None, None -> None
        | Some a, _, _, _
        | _, Some a, _, _
        | _, _, Some a, _
        | _, _, _, Some a -> Some a
      )
  and contains_rewrite_b (b: Types.b_exp) : Types.a_exp option =
    match b with
    | Boolean (_) -> None
    | BAnd (b1, b2)
    | BOr (b1, b2) -> (
      match contains_rewrite_b b1, contains_rewrite_b b2 with
        None, None -> None
      | Some a, _
      | _, Some a -> Some a
    )
    | BNot (b) -> contains_rewrite_b b
    | BCmp (a1, a2)
    | BCmpLess (a1, a2)
    | BCmpLessEq (a1, a2)
    | BCmpGreater (a1, a2)
    | BCmpGreaterEq (a1, a2) -> (
        match contains_rewrite_a a1, contains_rewrite_a a2 with
          None, None -> None
        | Some a, _
        | _, Some a -> Some a
      )
  and contains_rewrite_a (a: Types.a_exp) : Types.a_exp option =
    match a with
    | Variable _
    | Integer _ -> None
    | Plus (a1, a2)
    | Minus (a1, a2)
    | Times (a1, a2)
    | Division (a1, a2)
    | Modulo (a1, a2)
    | Power (a1, a2) -> (
        match contains_rewrite_a a1, contains_rewrite_a a2 with
          None, None -> None
        | Some a, _
        | _, Some a -> Some a
      )
    | PowerMod (_) -> Some a
    | Rand (a) -> contains_rewrite_a a
  in

  (* obtain the list of used variables so that fresh ones can be created *)
  let rec uv (prg: Types.c_exp) : Types.variable list =
    match prg with
    | Skip -> []
    | Assignment (x, _) -> [x]
    | Sequence (c1, c2) -> uv c1 @ uv c2
    | If (_, c1, c2) -> uv c1 @ uv c2
    | While (_, c) -> uv c
    | For (c1, _, c2, c3) -> uv c1 @ uv c2 @ uv c3
  in
  let usedvariables = i :: o :: (uv prg) in

  let counter = ref 0 in
  let new_fresh_var (pref: string) : Types.variable =
    let rec h () =
      let candidate = pref ^ (string_of_int !counter) in
      if (List.mem candidate usedvariables) then (
        counter := !counter + 1;
        h ()
      ) else (
        counter := !counter + 1;
        candidate
      )
    in
    h ()
  in

  (* functions that replace a pattern in a subexpression *)
  let rec replace_occurrence_a
      (pattern: Types.a_exp)
      (replace: Types.a_exp)
      (a: Types.a_exp)
    : Types.a_exp =
    if a = pattern then
      replace
    else (
      let r_o_a = replace_occurrence_a pattern replace in
      match a with
      | Variable _
      | Integer _ -> a
      | Plus (a1, a2) -> Plus (r_o_a a1, r_o_a a2)
      | Minus (a1, a2) -> Minus (r_o_a a1, r_o_a a2)
      | Times (a1, a2) -> Times (r_o_a a1, r_o_a a2)
      | Division (a1, a2) -> Division (r_o_a a1, r_o_a a2)
      | Modulo (a1, a2) -> Modulo (r_o_a a1, r_o_a a2)
      | Power (a1, a2) -> Power (r_o_a a1, r_o_a a2)
      | PowerMod (a1, a2, a3) -> PowerMod (r_o_a a1, r_o_a a2, r_o_a a3)
      | Rand (a) -> Rand (r_o_a a)
    )
  and replace_occurrence_b
      (pattern: Types.a_exp)
      (replace: Types.a_exp)
      (b: Types.b_exp)
    : Types.b_exp =
    let r_o_b = replace_occurrence_b pattern replace in
    let r_o_a = replace_occurrence_a pattern replace in
    match b with
    | Boolean _ -> b
    | BAnd (b1, b2) -> BAnd (r_o_b b1, r_o_b b2)
    | BOr (b1, b2) -> BOr (r_o_b b1, r_o_b b2)
    | BNot (b) -> BNot (r_o_b b)
    | BCmp (a1, a2) -> BCmp (r_o_a a1, r_o_a a2)
    | BCmpLess (a1, a2) -> BCmpLess (r_o_a a1, r_o_a a2)
    | BCmpLessEq (a1, a2) -> BCmpLessEq (r_o_a a1, r_o_a a2)
    | BCmpGreater (a1, a2) -> BCmpGreater (r_o_a a1, r_o_a a2)
    | BCmpGreaterEq (a1, a2) -> BCmpGreaterEq (r_o_a a1, r_o_a a2)
  in

  (* function that creates the equivalent code for a powermod using simpler
     instructions *)
  let partial freshres a1 a2 a3 : Types.c_exp =
    let freshpow = new_fresh_var "pow" in
    let freshexp = new_fresh_var "exp" in
    let freshmod = new_fresh_var "mod" in
    Sequence (
      Sequence (
        Sequence (
          Assignment (freshpow, a1),
          Assignment (freshexp, a2)),
        Sequence (
          Assignment (freshmod, a3),
          Assignment (freshres, Integer 1))),
      Sequence (
        If (BCmpLess (Variable freshexp, Integer 0),
            Assignment (freshexp, Minus (Integer 0, Variable freshexp)),
            Skip),
        While (
          BCmpGreater (Variable freshexp, Integer 0),
          Sequence (
            If (BCmp (Integer 1, Modulo (Variable freshexp, Integer 2)),
                Assignment (freshres,
                            Modulo (Times (Variable freshres,
                                           Variable freshpow),
                                    Variable freshmod)),
                Skip),
            Sequence (
              Assignment (freshpow,
                          Modulo (Times (Variable freshpow,
                                         Variable freshpow),
                                  Variable freshmod)),
              Assignment (freshexp, Division (Variable freshexp, Integer 2))
            )))))
  in

  let replace_pwm (pwm: Types.a_exp) (p: Types.c_exp) : Types.c_exp =
    match pwm, p with
    | PowerMod (a1, a2, a3), Assignment (x, a) ->
      let freshres = new_fresh_var "res" in
      Sequence (
        partial freshres a1 a2 a3,
        Assignment(x, replace_occurrence_a pwm (Variable freshres) a)
      )
    | PowerMod (a1, a2, a3), If (b, ifa1, ifa2) ->
      let freshres = new_fresh_var "res" in
      Sequence (
        partial freshres a1 a2 a3,
        If (replace_occurrence_b pwm (Variable freshres) b, ifa1, ifa2)
      )
    | PowerMod (a1, a2, a3), While (b, wa) ->
      let freshres = new_fresh_var "res" in
      Sequence (
        partial freshres a1 a2 a3,
        While (replace_occurrence_b pwm (Variable freshres) b, wa)
      )
    | PowerMod (a1, a2, a3), For (fora1, b, fora2, fora3) ->
      let freshres = new_fresh_var "res" in
      Sequence (
        partial freshres a1 a2 a3,
        For ( fora1,
              replace_occurrence_b pwm (Variable freshres) b,
              fora2,
              fora3 )
      )
    | _ -> failwith "PowerMod is not present"
  in

  let rec rw_a (prg: Types.c_exp) : Types.c_exp =
    match prg with
    | Skip -> Skip
    | Assignment (x, a) -> (
        match contains_rewrite_a a with
          None -> Assignment (x, a)
        | Some (PowerMod (a1, a2, a3)) -> (
            replace_pwm (PowerMod (a1, a2, a3)) prg
          )
        | Some _ -> failwith "Found powmod then lost it."
      )
    | Sequence (c1, c2) -> Sequence (rw_a c1, rw_a c2)
    | If (b, c1, c2) -> (
      match contains_rewrite_b b with
        None -> If (b, rw_a c1, rw_a c2)
      | Some (PowerMod (a1, a2, a3)) ->
        replace_pwm (PowerMod (a1, a2, a3)) prg
      | Some _ -> failwith "Found powmod then lost it."
    )
    | While (b, c) -> (
      match contains_rewrite_b b with
        None -> While (b, rw_a c)
      | Some (PowerMod (a1, a2, a3)) ->
        replace_pwm (PowerMod (a1, a2, a3)) prg
      | Some _ -> failwith "Found powmod then lost it."
    )
    | For (c1, b, c2, c3) -> (
      match contains_rewrite_b b with
        None -> For (rw_a c1, b, rw_a c2, rw_a c3)
      | Some (PowerMod (a1, a2, a3)) ->
        replace_pwm (PowerMod (a1, a2, a3)) prg
      | Some _ -> failwith "Found powmod then lost it."
    )
  in

  (* we first check that at least one powermod is present *)
  if Option.is_none (contains_rewrite prg) then
    Main (i, o, prg)
  else
    Main (i, o, rw_a prg)