fpOpt.lem

(*
  Definition of the fp_pattern language for Icing optimizations
*)

open import Pervasives
open import Lib
open import FpValTree

type fp_pat =
     | Word of word64
     | PatVar of nat
     | Unop of fp_uop * fp_pat
     | Binop of fp_bop * fp_pat * fp_pat
     | Terop of fp_top * fp_pat * fp_pat * fp_pat
     (* | Pred of fp_pred * fp_pat *)
     | Cmp of fp_cmp * fp_pat * fp_pat
     | Optimise of fp_opt * fp_pat

(* Substitutions are maps (paired lists) from numbers to 'a *)
type subst 'v = list (nat * 'v)

val substLookup: forall 'v. subst 'v -> nat -> maybe 'v
let rec substLookup ([]) n = Nothing
    and substLookup ((m, v)::s) n =
      if (m = n) then Just v else substLookup s n
declare termination_argument substLookup = automatic

val substUpdate: forall 'v. nat -> 'v -> subst 'v -> maybe (subst 'v)
let rec substUpdate n v1 s =
    match s with
    | [] -> Nothing
    | (s1::sN) ->
      match s1 with
      | (m,v2) ->
        if (n = m) then Just ((n,v1)::sN)
        else
          match (substUpdate n v1 sN) with
          | Nothing -> Nothing
          | Just sNew -> Just ((m,v2)::sNew)
          end
      end
    end

val substAdd: forall 'v. nat -> 'v -> subst 'v -> subst 'v
let rec substAdd n v s =
    match (substUpdate n v s) with
    | Just sNew -> sNew
    | Nothing -> (n,v)::s
    end

(* Matching a fp_pattern with the top-level of a value tree,
  if a matching exists an option with a substitution is returned.
  The matcher takes an additional substituion as argument to make sure
  that we do not double match a fp_pattern to different expressions
*)
val matchWordTree: fp_pat -> fp_word_val -> subst fp_word_val -> maybe (subst fp_word_val)
let rec matchWordTree (Word w1) (Fp_const w2) s =
    if (w1 = w2) then Just s else Nothing
    and matchWordTree (PatVar n) v s =
      match substLookup s n with
      | Just v1 -> if v1 = v then Just s else Nothing
      | Nothing -> Just (substAdd n v s)
      end
    and matchWordTree (Unop op1 p) (Fp_uop op2 v) s =
      if (op1 = op2)
      then matchWordTree p v s
      else Nothing
    and matchWordTree (Binop b1 p1 p2) (Fp_bop b2 v1 v2) s =
      if (b1 = b2)
      then
        match matchWordTree p1 v1 s with
        | Nothing -> Nothing
        | Just s1 -> matchWordTree p2 v2 s1
        end
      else Nothing
    and matchWordTree (Terop t1 p1 p2 p3) (Fp_top t2 v1 v2 v3) s =
      if (t1 = t2)
      then
        match matchWordTree p1 v1 s with
        | Nothing -> Nothing
        | Just s1 ->
          match matchWordTree p2 v2 s1 with
          | Nothing -> Nothing
          | Just s2 -> matchWordTree p3 v3 s2
          end
        end
      else Nothing
    and matchWordTree (Optimise fp_opt1 p) (Fp_wopt fp_opt2 v) s =
      if fp_opt1 = fp_opt2 then matchWordTree p v s else Nothing
    and matchWordTree _ _ s = Nothing

val matchBoolTree: fp_pat -> fp_bool_val -> subst fp_word_val -> maybe (subst fp_word_val)
let rec matchBoolTree (Optimise fp_opt1 p) (Fp_bopt fp_opt2 v) s =
      if fp_opt1 = fp_opt2 then (matchBoolTree p v s) else Nothing
    (* and matchBoolTree (Pred pred1 p) (Fp_pred pred2 v) s =
      if (pred1 = pred2) then matchWordTree p v s else Nothing *)
    and matchBoolTree (Cmp cmp1 p1 p2) (Fp_cmp cmp2 v1 v2) s =
      if (cmp1 = cmp2)
      then
        match matchWordTree p1 v1 s with
        | Nothing -> Nothing
        | Just s1 -> matchWordTree p2 v2 s1
        end
      else Nothing
    and matchBoolTree _ _ _ = Nothing

(* Instantiate a given fp_pattern with a substitution into a value tree *)
val instWordTree: fp_pat -> subst fp_word_val -> maybe fp_word_val
let rec instWordTree (Word w) s = Just (Fp_const w)
    and instWordTree (PatVar n) s = substLookup s n
    and instWordTree (Unop u p) s =
      match instWordTree p s with
      | Nothing -> Nothing
      | Just v -> Just (Fp_uop u v)
      end
    and instWordTree (Binop op p1 p2) s =
      match (instWordTree p1 s, instWordTree p2 s) with
      | (Just v1, Just v2) -> Just (Fp_bop op v1 v2)
      | (_, _) -> Nothing
      end
    and instWordTree (Terop op p1 p2 p3) s =
      match (instWordTree p1 s, instWordTree p2 s , instWordTree p3 s) with
      | (Just v1, Just v2, Just v3) -> Just (Fp_top op v1 v2 v3)
      | (_, _, _) -> Nothing
      end
    and instWordTree (Optimise fp_opt p) s =
      match instWordTree p s with
      | Nothing -> Nothing
      | Just v -> Just (Fp_wopt fp_opt v)
      end
    and instWordTree _ _ = Nothing

val instBoolTree: fp_pat -> subst fp_word_val -> maybe fp_bool_val
let rec instBoolTree (Optimise fp_opt p) s =
      match instBoolTree p s with
      | Nothing -> Nothing
      | Just v -> Just (Fp_bopt fp_opt v)
      end
    (* and instBoolTree (Pred pr p1) s =
      match instWordTree p1 s with
      | Nothing -> Nothing
      | Just v -> Just (Fp_pred pr v)
      end *)
    and instBoolTree (Cmp cmp p1 p2) s =
      match (instWordTree p1 s, instWordTree p2 s) with
      | (Just v1, Just v2) -> Just (Fp_cmp cmp v1 v2)
      | (_, _) -> Nothing
      end
    and instBoolTree _ _ = Nothing

(* Define a floating-point rewrite as a pair of a source and target fp_pattern *)
type fp_rw = (fp_pat * fp_pat)

(** Rewriting on value trees is done in the semantics by picking a fp_path
  that walks down the value tree structure and then applies the rewrite in place
  if it matches **)

(* Datatype for fp_paths into a value tree. Here is the leaf node meaning that the
  rewrite should be applied *)
type fp_path = | Left of fp_path | Right of fp_path | Center of fp_path | Here

val maybe_map: forall 'v1 'v2. ('v1 -> 'v2) -> maybe 'v1 -> maybe 'v2
let rec maybe_map f Nothing = Nothing
    and maybe_map f (Just res) = Just (f res)

(* Function rwFp_pathValTree b rw p v recurses through value tree v using fp_path p
  until p = Here or no further recursion is possible because of a mismatch.
  In case of a mismatch the function simply returns Nothing. *)
val rwFp_pathWordTree: fp_rw -> fp_path -> fp_word_val -> maybe fp_word_val
let rec rwFp_pathWordTree rw Here v =
      let (lhs, rhs) = rw in
      match matchWordTree lhs v [] with
        | Nothing -> Nothing
        | Just s -> instWordTree rhs s
      end
    and rwFp_pathWordTree rw (Left p) (Fp_bop op v1 v2) =
      maybe_map (fun v1 -> Fp_bop op v1 v2) (rwFp_pathWordTree rw p v1)
    and rwFp_pathWordTree rw (Right p) (Fp_bop op v1 v2) =
      maybe_map (fun v2 -> Fp_bop op v1 v2) (rwFp_pathWordTree rw p v2)
    and rwFp_pathWordTree rw (Center p) (Fp_uop op v1) =
      maybe_map (fun v -> Fp_uop op v) (rwFp_pathWordTree rw p v1)
    and rwFp_pathWordTree rw (Left p) (Fp_top op v1 v2 v3) =
      maybe_map (fun v1 -> Fp_top op v1 v2 v3) (rwFp_pathWordTree rw p v1)
    and rwFp_pathWordTree rw (Center p) (Fp_top op v1 v2 v3) =
      maybe_map (fun v2 -> Fp_top op v1 v2 v3) (rwFp_pathWordTree rw p v2)
    and rwFp_pathWordTree rw (Right p) (Fp_top op v1 v2 v3) =
      maybe_map (fun v3 -> Fp_top op v1 v2 v3) (rwFp_pathWordTree rw p v3)
    and rwFp_pathWordTree rw (Center p) (Fp_wopt fp_opt v) =
      maybe_map (fun v -> Fp_wopt fp_opt v) (rwFp_pathWordTree rw p v)
    and rwFp_pathWordTree _ _ _ = Nothing

val rwFp_pathBoolTree: fp_rw -> fp_path -> fp_bool_val -> maybe fp_bool_val
let rec rwFp_pathBoolTree rw Here v =
      let (lhs, rhs) = rw in
      match matchBoolTree lhs v [] with
        | Nothing -> Nothing
        | Just s -> instBoolTree rhs s
      end
    and rwFp_pathBoolTree rw (Center p) (Fp_bopt fp_opt v) =
      maybe_map (fun v -> Fp_bopt fp_opt v) (rwFp_pathBoolTree rw p v)
    (* and rwFp_pathBoolTree rw (Center p) (Fp_pred pr v) =
      maybe_map (fun v -> Fp_pred pr v) (rwFp_pathWordTree rw p v) *)
    and rwFp_pathBoolTree rw (Left p) (Fp_cmp cmp v1 v2) =
      maybe_map (fun v1 -> Fp_cmp cmp v1 v2) (rwFp_pathWordTree rw p v1)
    and rwFp_pathBoolTree rw (Right p) (Fp_cmp cmp v1 v2) =
      maybe_map (fun v2 -> Fp_cmp cmp v1 v2) (rwFp_pathWordTree rw p v2)
    and rwFp_pathBoolTree _ _ _ = Nothing

(* Datatype holding a single rewrite application in the form of a fp_path into the
  value tree and a number giving the index of the rewrite to be used *)
type rewrite_app = | RewriteApp of fp_path * nat (* which rewrite rule *)

val nth: forall 'v. list 'v -> nat -> maybe 'v
let rec nth [] n = Nothing
    and nth (x::xs) n =
      if (n = 0) then Nothing
      else if (n = 1) then Just x
      else nth xs (n-1)

(* rwAllValTree rwApps canOpt rws v applies all the rewrite_app's in rwApps to
    value tree v using rwFp_pathValTree *)
val rwAllWordTree: list rewrite_app -> list fp_rw -> fp_word_val -> maybe fp_word_val
let rec rwAllWordTree [] rws v = Just v
    and rwAllWordTree ((RewriteApp p n)::rs) rws v =
      match nth rws n with
      | Nothing -> Nothing
      | Just rw ->
        match rwFp_pathWordTree rw p v with
        | Nothing -> Nothing
        | Just vNew -> rwAllWordTree rs rws vNew
        end
      end

val rwAllBoolTree: list rewrite_app -> list fp_rw -> fp_bool_val -> maybe fp_bool_val
let rec rwAllBoolTree [] rws v = Just v
    and rwAllBoolTree ((RewriteApp p n)::rs) rws v =
      match nth rws n with
      | Nothing -> Nothing
      | Just rw ->
        match rwFp_pathBoolTree rw p v with
        | Nothing -> Nothing
        | Just vNew -> rwAllBoolTree rs rws vNew
        end
      end

val no_fp_opts : nat -> list rewrite_app
let no_fp_opts (n:nat) = []