transform small-step normalizer in direct style to big-step directly

refocus_direct.ml

type op = Add | Sub

type expr =
  | Lit of int
  | Op  of expr * op * expr

type result =
  | Val of int
  | Term of expr
  | Stuck

(* plain small-step reduction implemented with structural recursion *)
let rec reduce expr =
  match expr with
  | Lit i -> Val i
  | Op (l, op, r) -> 
      match reduce l with
      | Val x ->
          (match reduce r with
            | Val y ->
                (match op with
                  | Add -> Term (Lit (x + y))
                  | Sub ->
                      if x < y
                      then Stuck
                      else Term (Lit (x - y)))
            | Term r' -> Term (Op (Lit x, op, r'))
            | Stuck -> Stuck)
      | Term l' -> Term (Op (l', op, r))
      | Stuck -> Stuck

let rec original_normalize expr =
  match reduce expr with
  | Val v -> Some v
  | Stuck -> None
  | Term t -> original_normalize t

(* inline [reduce] into [normalize] and do case-of-case *)
module V1 = struct
  let rec normalize expr =
    match expr with
    | Lit i -> Some i
    | Op (l, op, r) ->
        match reduce l with
        | Val x ->
            (match reduce r with
              | Val y ->
                  (match op with
                    | Add -> Some (x + y)
                    | Sub ->
                        if x < y
                        then None
                        else Some (x - y))
              | Term r' -> normalize (Op (Lit x, op, r'))
              | Stuck -> None)
        | Term l' -> normalize (Op (l', op, r))
        | Stuck -> None
end

(* move the handling of [reduce l] and [reduce r] to a separate function *)
module V2 = struct
  let rec normalize expr =
    match expr with
    | Lit i -> Some i
    | Op (l, op, r) -> normalize_l l op r

  and normalize_l l op r =
    match reduce l with
    | Val x -> normalize_r x op r
    | Term l' -> normalize (Op (l', op, r))
    | Stuck -> None

  and normalize_r x op r =
    match reduce r with
    | Val y ->
        (match op with
          | Add -> Some (x + y)
          | Sub ->
              if x < y
              then None
              else Some (x - y))
    | Term r' -> normalize (Op (Lit x, op, r'))
    | Stuck -> None
end

(* inline the recursive call to [normalize] in [normalize_l] and [normalize_r] *)
module V3 = struct
  let rec normalize expr =
    match expr with
    | Lit i -> Some i
    | Op (l, op, r) -> normalize_l l op r

  and normalize_l l op r =
    match reduce l with
    | Val x -> normalize_r x op r
    | Term l' ->
        (* initial
           = normalize (Op (l', op, r))
        *) normalize_l l' op r
    | Stuck -> None

  and normalize_r x op r =
    match reduce r with
    | Val y ->
        (match op with
          | Add -> Some (x + y)
          | Sub ->
              if x < y
              then None
              else Some (x - y))
    | Term r' ->
        (* initial
           = normalize (Op (Lit x, op, r'))
           = normalize_l (Lit x) op r'
        *) normalize_r x op r'
    | Stuck -> None
end

(* move the base case (value case) of [normalize_l] and [normalize_r] out *)
module V4 = struct
  let rec normalize expr =
    match expr with
    | Lit i -> Some i
    | Op (l, op, r) ->
        (match normalize_l l with
          | Some x ->
              (* originally the [Val] case of [normalize_l] *)
              (match normalize_r r with
                | Some y ->
                    (* originally the [Val] case of [normalize_r] *)
                    (match op with
                      | Add -> Some (x + y)
                      | Sub ->
                          if x < y
                          then None
                          else Some (x - y))
                | None -> None)
          | None -> None)

  and normalize_l l =
    match reduce l with
    | Val x -> Some x
    | Term l' -> normalize_l l'
    | Stuck -> None

  and normalize_r r =
    match reduce r with
    | Val y -> Some y
    | Term r' -> normalize_r r'
    | Stuck -> None
end

(* Finally, notice that [normalize_l] and [normalize_r] are exactly the same as
   [original_normalize], so we can replace [normalize_l] and [normalize_r] with [original_normalize]!

   We also know that [original_normalize] is the same as [normalize],
   because [normalize] is obtained by transformation on [original_normalize],
   so we can further replace [normalize_l] and [normalize_r] with [normalize]!

   This give us a direct style, big step evaluator!
*)
module V5 = struct
  let rec normalize expr =
    match expr with
    | Lit i -> Some i
    | Op (l, op, r) ->
        (match normalize l with
          | Some x ->
              (match normalize r with
                | Some y ->
                    (match op with
                      | Add -> Some (x + y)
                      | Sub ->
                          if x < y
                          then None
                          else Some (x - y))
                | None -> None)
          | None -> None)
end


(* ======= another example ====== *)

(* lambda term with de bruijn index *)
type term =
  | Var of int
  | Lam of term
  | App of term * term

(* shift all index beyond [base] in a term up by one *)
let rec shift_up ~base t =
  match t with
  | Var i      -> Var (if i < base then i else i + 1)
  | Lam t'     -> Lam (shift_up ~base:(base + 1) t')
  | App (f, a) -> App (shift_up ~base f, shift_up ~base a)

(* perform [shift_up] for [shift] times *)
let rec original_shift_up_n ~base t ~shift =
  if shift = 0
  then t
  else original_shift_up_n ~base (shift_up ~base t) ~shift:(shift - 1)

(* step 1: inline [shift_up] in [shift_up_n] and do case-of-case transformation *)
module V1 = struct
  let rec shift_up_n ~base t ~shift =
    if shift = 0
    then t
    else
      match t with
      | Var i -> shift_up_n ~base (Var (if i < base then i else i + 1)) ~shift:(shift - 1)
      | Lam t' -> shift_up_n ~base (Lam (shift_up ~base:(base + 1) t')) ~shift:(shift - 1)
      | App (f, a) ->
          shift_up_n ~base (App (shift_up ~base f, shift_up ~base a)) ~shift:(shift - 1)
end

(* step 2: split the three cases of [t] out into separate functions *)
module V2 = struct
  let rec shift_up_n ~base t ~shift =
    if shift = 0
    then t
    else
      match t with
      | Var i -> shift_up_n_var_case ~base i ~shift
      | Lam t' -> shift_up_n_lam_case ~base t' ~shift
      | App (f, a) -> shift_up_n_app_case ~base f a ~shift

  (* [shift_up_n_var_case ~base i ~shift] = [shift_up_n ~base (Var i) ~shift] *)
  and shift_up_n_var_case ~base i ~shift =
    if shift = 0
    then Var i
    else
      (* originally:

           shift_up_n ~base (Var (if i < base then i else i + 1)) ~shift:(shift - 1)

      *) shift_up_n_var_case ~base (if i < base then i else i + 1) ~shift:(shift - 1)

  (* [shift_up_n_lam_case ~base i ~shift] = [shift_up_n ~base (Lam t) ~shift] *)
  and shift_up_n_lam_case ~base t ~shift =
    if shift = 0
    then Lam t
    else
      (* originally:

          shift_up_n ~base (Lam (shift_up ~base:(base + 1) t)) ~shift:(shift - 1)
      *)
      shift_up_n_lam_case ~base (shift_up ~base:(base + 1) t) ~shift:(shift - 1)

  (* [shift_up_n_app_case ~base f a ~shift] = [shift_up_n ~base (App (f, a)) ~shift] *)
  and shift_up_n_app_case ~base f a ~shift =
    if shift = 0
    then App (f, a)
    else
      (* originally:

          shift_up_n ~base (App (shift_up ~base f, shift_up ~base a)) ~shift:(shift - 1)
      *)
      shift_up_n_app_case ~base (shift_up ~base f) (shift_up ~base a) ~shift:(shift - 1)
end

(* step 3:

   - lift the base case of the recursive function for var/lam/app cases out
   - and notice that, computation of [f] and [a] never depend on each other in the app case.
     So we can split [shift_up_n_app_case] into two independent functions computing [f] and [a]
*)
module V3 = struct
  let rec shift_up_n ~base t ~shift =
    if shift = 0
    then t
    else
      match t with
      | Var i -> Var (shift_up_n_var_case ~base i ~shift)
      | Lam t' -> Lam (shift_up_n_lam_case ~base t' ~shift)
      | App (f, a) ->
          App
            ( shift_up_n_app_case_f_part ~base f ~shift
            , shift_up_n_app_case_a_part ~base a ~shift
            )

  and shift_up_n_var_case ~base i ~shift =
    if shift = 0
    then i
    else shift_up_n_var_case ~base (if i < base then i else i + 1) ~shift:(shift - 1)

  and shift_up_n_lam_case ~base t ~shift =
    if shift = 0
    then t
    else shift_up_n_lam_case ~base (shift_up ~base:(base + 1) t) ~shift:(shift - 1)

  and shift_up_n_app_case_f_part ~base f ~shift =
    if shift = 0
    then f
    else shift_up_n_app_case_f_part ~base (shift_up ~base f) ~shift:(shift - 1)

  and shift_up_n_app_case_a_part ~base a ~shift =
    if shift = 0
    then a
    else shift_up_n_app_case_a_part ~base (shift_up ~base a) ~shift:(shift - 1)
end

(* step 4:

   now here comes the final eureka moment: observe that:

   - [shift_up_n_lam_case ~base] is exactly [original_shift_up_n ~base:(base + 1)]
   - [shift_up_n_app_case_f_part] and [shift_up_n_app_case_a_part]
     are exactly [original_shift_up_n]

   so we can replace all these three with [original_shift_up_n].
   Since [original_shift_up_n] is the same as [shift_up_n],
   we can further replace them wih [shift_up_n].
*)
module V4 = struct
  let rec shift_up_n ~base t ~shift =
    if shift = 0
    then t
    else
      match t with
      | Var i      -> Var (shift_up_n_var_case ~base i ~shift)
      | Lam t'     -> Lam (shift_up_n ~base:(base + 1) t' ~shift)
      | App (f, a) -> App (shift_up_n ~base f ~shift, shift_up_n ~base a ~shift)

  and shift_up_n_var_case ~base i ~shift =
    if shift = 0
    then i
    else shift_up_n_var_case ~base (if i < base then i else i + 1) ~shift:(shift - 1)
end

(* step 5:

   Finally, we can simplify [shift_up_n_var_case] with some ad-hoc observation.
   Observe that the value of [i < base] never change in recursive call:

   - if [i < base], the next value of [i] will not change
   - if [i >= base], the next value of [i] will be [i + 1 > i >= base]

   so we can split [shift_up_n_var_case] into two functions, for [i < base] and [i >= base]:

   - in the [i < base] case, the value of [i] never change, so the result if just [i]
   - the case of [i >= base] looks like this:

       let rec var_case_i_ge_base_case ~base i ~shift =
         if shift = 0
         then i
         else var_case_i_ge_base_case ~base (i + 1) ~shift:(shift - 1)

     which is exactly [i + shift]

   so we can simplify [shift_up_n_var_case] into:

     if i < base then i else i + shift

   and we inline this definition into [shift_up_n]
*)
module V5 = struct
  let rec shift_up_n ~base t ~shift =
    if shift = 0
    then t
    else
      match t with
      | Var i      -> Var (if i < base then i else i + shift)
      | Lam t'     -> Lam (shift_up_n ~base:(base + 1) t' ~shift)
      | App (f, a) -> App (shift_up_n ~base f ~shift, shift_up_n ~base a ~shift)
end