lazy-optimal.ml

type index = int

type term =
  | T_var of index
  | T_lambda of term lazy_t
  | T_apply of term lazy_t * term lazy_t

let rec shift ~by ~depth term =
  match term with
  | T_var var -> (
      match var <= depth with true -> T_var (by + var) | false -> T_var var)
  | T_lambda body ->
      let body =
        let depth = 1 + depth in
        shift_lazy ~by ~depth body
      in
      T_lambda body
  | T_apply (funct, arg) ->
      let funct = shift_lazy ~by ~depth funct in
      let arg = shift_lazy ~by ~depth arg in
      T_apply (funct, arg)

and shift_lazy ~by ~depth term = lazy (shift ~by ~depth (Lazy.force term))

let shift ~by term = shift ~by ~depth:0 term

let rec subst ~to_ ~depth term =
  match term with
  | T_var var -> (
      let to_ = Lazy.force to_ in
      match var = depth with true -> shift ~by:depth to_ | false -> T_var var)
  | T_lambda body ->
      let body =
        let depth = 1 + depth in
        subst_lazy ~to_ ~depth body
      in
      T_lambda body
  | T_apply (funct, arg) ->
      let funct = subst_lazy ~to_ ~depth funct in
      let arg = subst_lazy ~to_ ~depth arg in
      T_apply (funct, arg)

and subst_lazy ~to_ ~depth term = lazy (subst ~to_ ~depth (Lazy.force term))

let subst ~to_ term = subst ~to_ ~depth:0 term

let rec eval term =
  match term with
  | T_var var -> T_var var
  | T_lambda body ->
      let body = eval_lazy body in
      T_lambda body
  | T_apply (funct, arg) -> (
      let funct = eval_lazy funct in
      let arg = eval_lazy arg in
      match Lazy.force funct with
      | T_lambda body ->
          let body = Lazy.force body in
          eval @@ subst ~to_:arg body
      | T_var _ | T_apply (_, _) -> T_apply (funct, arg))

and eval_lazy term = lazy (eval (Lazy.force term))