Skip to content

Instantly share code, notes, and snippets.

@Guest0x0

Guest0x0/refocus_generic.ml

Created January 7, 2025 17:35
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save Guest0x0/f6f7f55e59c26ac84793117c539c3c8b to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save Guest0x0/f6f7f55e59c26ac84793117c539c3c8b to your computer and use it in GitHub Desktop.
Download ZIP
a generic formalization & evaluation of the "refocus" technique, connecting small-step and big-step evaluators
Raw
refocus_generic.ml
module type Reduction_Theory = sig
type term
type value
type redex
val v2t : value -> term
val contract : redex -> term option
(* an elementary context is a single frame of elemetary context,
for example:
- for CBN lambda calculus, `elem_ctx` will be just `[] t`
- for CBV lambda calculus, `elem_ctx` will be `[] t` or `v []`
*)
type elem_ctx
(* plug a term into an elementary context *)
val plug : elem_ctx -> term -> term
(* This is the core construct of our abstraction of reduction theory:
a single-step decomposition operation.
The three kinds of result are:
- `Value v`, indicating the term is already a value
- `Redex r`, indicating the term is already a redex
- `Decomp c t`, indicating that:
1. original term can be decomposed into elementary context `c` and sub term `t`
2. `t` is the next sub term to evaluate
*)
type decomp_result =
| Value of value
| Redex of redex
| Decomp of elem_ctx * term
val one_step_decomp : term -> decomp_result
(* The operations above should satisfy several properties *)
(* 1. sub term generated by [one_step_decomp] is never value:
one_step_decomp t = Decomp (c, t') -> `t'` is not value
This is the most crucial property,
also the main difference with more common approaches.
In usual implementation, decomposition will blindly dive into the first sub-term.
If the sub-term is already a value, it backtrack to the caller and seek the next sub-term,
until a non-value sub term is found, or the expression turn out to be already a redex.
In the formalization here, instead of try-and-backtrack,
we require `one_step_decomp` to inspect its sub-term
and select the first non-value sub-term for us.
*)
(* 2. `one_step_decomp` is left inverse of `plug` for non-values:
`t` is not value -> one_step_decomp (plug c t) = Decomp (c, t)
It is crucial that this property does *NOT* hold for the value case.
`one_step_decomp (plug c v)` should return the next sub-term to evaluate or return a redex directly.
*)
(* 3. evaluable forms are disjoint with values:
`plug c t` is not value for all `c` and `t`
*)
(* 4. `v2t` and `one_step_decomp` respect each other:
t = v2t v <-> one_step_decomp t = Value v
*)
end
(* a generic, small step normalizer in direct style *)
module V1 (R : Reduction_Theory) = struct
type result =
| V of R.value
| T of R.term
| Err
let rec reduce t =
match R.one_step_decomp t with
| Value v -> V v
| Redex r ->
(match R.contract r with
| Some t' -> T t'
| None -> Err)
| Decomp (ctx, t') ->
(match reduce t' with
| V v ->
(* since decomposed sub-term is never value,
the value case never occur in recursive call.
So this branch is unreachable
*)
assert false
| T t -> T (R.plug ctx t)
| Err -> Err)
let rec normalize t =
match reduce t with
| V v -> Some v
| T t -> normalize t
| Err -> None
end
(* CPS *)
module V2 (R : Reduction_Theory) = struct
type result =
| V of R.value
| T of R.term
| Err
let rec reduce t (k : result -> result) =
match R.one_step_decomp t with
| Value v -> k (V v)
| Redex r ->
(match R.contract r with
| Some t -> k (T t)
| None -> k Err)
| Decomp (ctx, t') ->
reduce t' (function
| V v -> assert false
| T t -> T (R.plug ctx t)
| Err -> Err)
let rec normalize t =
match reduce t (fun x -> x) with
| V v -> Some v
| T t -> normalize t
| Err -> None
end
(* split the error continuation out *)
module V3 (R : Reduction_Theory) = struct
type result =
| V of R.value
| T of R.term
| Err
type vot = V of R.value | T of R.term
let rec reduce t (k : vot -> result) =
match R.one_step_decomp t with
| Value v -> k (V v)
| Redex r ->
(match R.contract r with
| Some t -> k (T t)
| None -> Err)
| Decomp (ctx, t') ->
reduce t' (function
| V v -> assert false
| T t -> T (R.plug ctx t))
let rec normalize t =
match reduce t (function V v -> V v | T t -> T t) with
| V v -> Some v
| T t -> normalize t
| Err -> None
end
(* defunctionalize continuation *)
module V4 (R : Reduction_Theory) = struct
type result =
| V of R.value
| T of R.term
| Err
type vot = V of R.value | T of R.term
let rec reduce t (k : R.elem_ctx list) =
match R.one_step_decomp t with
| Value v -> continue_v k v
| Redex r ->
(match R.contract r with
| Some t -> continue_t k t
| None -> Err)
| Decomp (ctx, t') -> reduce t' (ctx :: k)
(* [continue_v k v = k (V v)] *)
and continue_v (k : R.elem_ctx list) (v : R.value) : result =
match k with
| [] -> V v
| ctx :: k' -> assert false
(* [continue_t k t = k (T t)] *)
and continue_t (k : R.elem_ctx list) (t : R.term) : result =
match k with
| [] -> T t
| ctx :: k' -> continue_t k' (R.plug ctx t)
let rec normalize t =
match reduce t [] with
| V v -> Some v
| T t -> normalize t
| Err -> None
end
(* 1. rewrite [normalize], lift [reduce] outside the main loop
2. [continue_t] always return [T], so we move the [T] constructor to its call site,
and rename [continue_t] to [recompose]
*)
module V5_1 (R : Reduction_Theory) = struct
type result =
| V of R.value
| T of R.term
| Err
type vot = V of R.value | T of R.term
(* originally [continue_t] *)
let rec recompose (k : R.elem_ctx list) (t : R.term) =
match k with
| [] -> t
| ctx :: k' -> recompose k' (R.plug ctx t)
let rec reduce t (k : R.elem_ctx list) =
match R.one_step_decomp t with
| Value v -> continue_v k v
| Redex r ->
(match R.contract r with
| Some t -> T (recompose k t)
| None -> Err)
| Decomp (ctx, t') -> reduce t' (ctx :: k)
and continue_v (k : R.elem_ctx list) (v : R.value) : result =
match k with
| [] -> V v
| ctx :: k' -> assert false
let rec normalize_loop (result : result) =
match result with
| V v -> Some v
| T t -> normalize_loop (reduce t [])
| Err -> None
let normalize t = normalize_loop (reduce t [])
end
(* do some fixedpoint promotion/demotion to expose decompose/contract/recompose structure:
[reduce] and [continue_v] together only have two base cases:
* the toplevel value case [V v]
* the redex case
we perform lightweight fixedpoint demotion to move the redex handling code outside,
and let [reduce] and [continue_v] return a new type [decomp_result].
Finally, rename [reduce] to [decompose], and now we get the decompose/contract/recompose
structure of reduction semantic
*)
module V5_2 (R : Reduction_Theory) = struct
type decomp_result =
| Decomp_value of R.value
| Decomp_redex of R.redex * R.elem_ctx list
let rec recompose (k : R.elem_ctx list) (t : R.term) : R.term =
match k with
| [] -> t
| ctx :: k' -> recompose k' (R.plug ctx t)
(* originally [reduce] *)
let rec decompose t (k : R.elem_ctx list) : decomp_result =
match R.one_step_decomp t with
| Value v -> continue_v k v
| Redex r ->
(* ===== CHANGES HERE ===== *)
Decomp_redex (r, k)
| Decomp (ctx, t') -> decompose t' (ctx :: k)
and continue_v (k : R.elem_ctx list) (v : R.value) =
match k with
| [] ->
(* ===== CHANGES HERE ===== *)
Decomp_value v
| ctx :: k' -> assert false
and normalize_loop (decomp : decomp_result) =
match decomp with
| Decomp_value v -> Some v
| Decomp_redex (r, k) ->
(* ===== CHANGES HERE =====
here we perform case-of-case transformation to simplify the code,
originally it should be:
match
(match R.contract r with
| Some t -> T (recompose k t)
| None -> Err)
with
| V v -> Some v
| T t -> normalize_loop (decompose t [])
| Err -> None
*)
(match R.contract r with
| Some t -> normalize_loop (decompose (recompose k t) [])
| None -> None)
let normalize t = normalize_loop (decompose t [])
end
(* now we can see the [decompose (recompose k t)] structure,
which is inefficient.
we want to introduce a new function [refocus t k]
that is equivant to [decompose (recompose k t) []].
Unlike existing refocus literature, in the setting here,
[refocus] is highly similar to, but not exactly the same as [decompose].
In particular, [refocus] is the same as [decompose] for all **non-value** input,
but differ from [decompose] in the value case.
In [decompose], the value case is unreachable with non-empty continuation.
However, in [refocus], the value + non-empty continuation case may happen.
Our [refocus] function is the same as usual presentation, and the difference lies in [decompose].
As mentioned in the comments in [Reduction_Theory],
[decompose] here does not perform back tracking on value case,
and this backtracking turns out to be exactly the extra bit of work [refocus] need to do!
That explains why [refocus] and [decompose] coindicide in usual presentation.
*)
module V6 (R : Reduction_Theory) = struct
type decomp_result =
| Decomp_value of R.value
| Decomp_redex of R.redex * R.elem_ctx list
let rec refocus t (k : R.elem_ctx list) : decomp_result =
match R.one_step_decomp t with
| Value v -> continue_v k v
| Redex r -> Decomp_redex (r, k)
| Decomp (ctx, t') -> refocus t' (ctx :: k)
and continue_v (k : R.elem_ctx list) (v : R.value) =
match k with
| [] -> Decomp_value v
| ctx :: k' ->
(* It is necessary to plug one layer before continuating [refocus] here.
Otherwise we will get a trivial infinite loop.
The intuition here of [refocus (plug ctx v) k'] here
is to seek the next sub-term to evaluate.
Instead of going all the way up with [recompose] and then dive down again,
we only plug one layer of elementary context here.
Since [plug ctx v] is never value, we know that [refocus (plug ctx v) k']
will always give us a redex or the next sub-term.
*)
refocus (R.plug ctx (R.v2t v)) k'
and normalize_loop (decomp : decomp_result) =
match decomp with
| Decomp_value v -> Some v
| Decomp_redex (r, k) ->
(match R.contract r with
| Some t -> normalize_loop (refocus t k)
| None -> None)
let normalize t = normalize_loop (refocus t [])
end
(* Now we give a paper-and-pen proof of the refocus theorem.
A formalized proof can be found in the Coq file.
Theorem refocus:
refocus t k = decompose (recompose k t) []
To prove the theorem, we need to first prove several lemmas:
Lemma 1: `refocus` is the same as `decompose` for non-value:
forall t k, `t` is not value -> refocus t k = decompose t k
Proof: simple induction.
Lemma 2: non-value part of refocus theorem:
forall t, `t` is not value -> refocus t k = decompose (recompose k t) []
Proof: induction on `k`:
(1). [k = []]:
right
= decompose t []
= refocus t [] (Lemma 1)
= left
(2). [k = ctx :: k']:
since `t` is not value, we have:
right
= decompose (recompose k' (plug ctx t))
= refocus (plug ctx t) k' (IH, `plug` is never value)
= refocus t (ctx :: k') (`one_step_decomp` is inverse of `plug` for non-value)
= left
Lemma 3: value part of refocus theorem:
forall v, refocus (v2t v) k = decompose (recompose k (v2t v)) []
Proof: induction on `k`:
(1). [k = []]: immediate
(2). [k = ctx :: k']:
right
= decompose (recompose k' (plug ctx (v2t v))) []
= refocus (plug ctx (v2t v)) k' (Lemma 2)
= left
Theorem refocus:
refocus t k = decompose (recompose k t) []
Proof: immediate from Lemma 2 and Lemma 3
Notice that the value case still only occur at toplevel in [refocus].
The difference with [decompose] is that the continuation may now be non-empty at toplevel.
This explains two things:
1. why [refocus] terminates in the value case
2. why we need to prove Lemma 2 first, and uses Lemma 2 in the proof of Lemma 3
*)
(* Now, we finish the last steps of turning the refocus'ed implementation
into a bit-step, direct style evaluator
*)
(* do lightweight fixedpoint promotion: fuse [normalize_loop] into [refocus] *)
module V7 (R : Reduction_Theory) = struct
(* originally [refocus] *)
let rec normalize_aux t (k : R.elem_ctx list) =
match R.one_step_decomp t with
| Value v -> continue_v k v
| Redex r ->
(match R.contract r with
| Some t -> normalize_aux t k
| None -> None)
| Decomp (ctx, t') -> normalize_aux t' (ctx :: k)
and continue_v (k : R.elem_ctx list) (v : R.value) =
match k with
| [] -> Some v
| ctx :: k' -> normalize_aux (R.plug ctx (R.v2t v)) k'
let normalize t = normalize_aux t []
end
(* refunctionalize the continuation *)
module V8 (R : Reduction_Theory) = struct
let rec normalize_aux t k =
match R.one_step_decomp t with
| Value v -> k v
| Redex r ->
(match R.contract r with
| Some t -> normalize_aux t k
| None -> None)
| Decomp (ctx, t') -> normalize_aux t' (fun v -> normalize_aux (R.plug ctx (R.v2t v)) k)
let normalize t = normalize_aux t (fun v -> Some v)
end
(* back to direct style *)
module V9 (R : Reduction_Theory) = struct
exception Stuck
let rec normalize_aux t =
match R.one_step_decomp t with
| Value v -> v
| Redex r ->
(match R.contract r with
| Some t -> normalize_aux t
| None -> raise Stuck)
| Decomp (ctx, t') ->
let v = normalize_aux t' in
normalize_aux (R.plug ctx (R.v2t v))
let normalize t = try Some (normalize_aux t) with Stuck -> None
end
(* example 1: simple arithmetic expression *)
module Arith = struct
module Theory = struct
type op = Add | Sub
type term =
| Lit of int
| Op of term * op * term
type value = int
let v2t v = Lit v
type redex = value * op * value
let contract (x, op, y) =
match op with
| Add -> Some (Lit (x + y))
| Sub ->
if x < y
then None
else Some (Lit (x - y))
type elem_ctx =
| L of op * term
| R of value * op
let plug ctx t =
match ctx with
| L (op, r) -> Op (t , op, r)
| R (x, op) -> Op (Lit x, op, t)
type decomp_result =
| Value of value
| Redex of redex
| Decomp of elem_ctx * term
let one_step_decomp t =
match t with
| Lit i -> Value i
| Op (Lit x, op, Lit y) -> Redex (x, op, y)
| Op (Lit x, op, r ) -> Decomp (R (x, op), r)
| Op (l , op, r ) -> Decomp (L (op, r), l)
end
module Reduce = V1(Theory)
module Eval = V9(Theory)
open Theory
let lit i = Lit i
let (<+>) x y = Op (x, Add, y)
let (<->) x y = Op (x, Sub, y)
let test expr ~expected =
let r1 = Reduce.normalize expr in
assert (r1 = expected);
let r2 = Eval.normalize expr in
assert (r2 = expected)
;;
let run_tests () =
test ~expected:(Some 42) ((lit 17 <+> lit 4) <+> (lit 25 <-> lit 4));
test ~expected:None (lit 1 <+> lit 2 <+> (lit 2 <-> lit 3))
;;
end
let () = Arith.run_tests ()
(* example 2:
call-by-value lambda calculus (de bruijn index + explicit substitution)
+ natural number
*)
module Lambda = struct
module Theory = struct
(* term without explicit substitution *)
type term0 =
| Lit0 of int
| Succ0
| Var0 of int
| Lam0 of term0
| App0 of term0 * term0
type term =
| Lit of int
| Succ
| App of term * term
| Sub of term0 * subst
and value =
| V_Lit of int
| V_Succ
| V_Clo of term0 * subst
and subst = value list
let v2t v =
match v with
| V_Lit i -> Lit i
| V_Succ -> Succ
| V_Clo (t, sub) -> Sub (Lam0 t, sub)
type elem_ctx =
| App_L of term
| App_R of value
let plug ctx t =
match ctx with
| App_L arg -> App (t, arg)
| App_R fn -> App (v2t fn, t)
type redex =
| R_Sub_Var of int * subst
| R_Sub_App of term0 * term0 * subst
| R_Beta of value * value
let contract r =
match r with
| R_Sub_Var (i, sub) -> List.nth_opt sub i |> Option.map v2t
| R_Sub_App (f, a, sub) -> Some (App (Sub (f, sub), Sub (a, sub)))
| R_Beta (V_Succ , V_Lit i) -> Some (Lit (i + 1))
| R_Beta (V_Clo (t, env), v) -> Some (Sub (t, v :: env))
| R_Beta _ -> None
type decomp_result =
| Value of value
| Redex of redex
| Decomp of elem_ctx * term
let as_value t : value option =
match t with
| Lit i -> Some (V_Lit i)
| Succ -> Some (V_Succ)
| Sub (Lam0 t, sub) -> Some (V_Clo (t, sub))
| Sub (Lit0 i, sub) -> Some (V_Lit i)
| Sub (Succ0 , sub) -> Some V_Succ
| _ -> None
let one_step_decomp t =
match t with
| Lit i -> Value (V_Lit i)
| Succ -> Value (V_Succ)
| Sub (Lit0 i , sub) -> Value (V_Lit i)
| Sub (Succ0 , sub) -> Value V_Succ
| Sub (Var0 x , sub) -> Redex (R_Sub_Var (x, sub))
| Sub (App0 (f, a), sub) -> Redex (R_Sub_App (f, a, sub))
| Sub (Lam0 t , sub) -> Value (V_Clo (t, sub))
| App (f, a) ->
(match as_value f with
| Some f ->
(match as_value a with
| Some a -> Redex (R_Beta (f, a))
| None -> Decomp (App_R f, a))
| None -> Decomp (App_L a, f))
end
module Reduce = V1(Theory)
module Eval = V9(Theory)
open Theory
let lit n = fun level -> Lit0 n
let succ = fun level -> Succ0
let lam t = fun level -> Lam0 (t (fun level' -> Var0 (level' - level - 1)) (level + 1))
let app f a = fun level -> App0 (f level, a level)
let app2 f x y = app (app f x) y
let mk_term pre_term = Sub (pre_term 0, [])
let church n =
let rec church_aux succ zero n =
match n with
| 0 -> zero
| n -> app succ (church_aux succ zero (n - 1))
in
lam @@ fun succ -> lam @@ fun zero -> church_aux succ zero n
let church_add =
lam @@ fun x ->
lam @@ fun y ->
lam @@ fun succ ->
lam @@ fun zero ->
app2 x succ (app2 y succ zero)
let church_mul =
lam @@ fun x ->
lam @@ fun y ->
lam @@ fun succ ->
lam @@ fun zero ->
app2 x (lam @@ fun acc -> app2 y succ acc) zero
let test pre_term ~expected =
let t = mk_term pre_term in
let r1 = Reduce.normalize t in
assert (r1 = expected);
let r2 = Eval.normalize t in
assert (r2 = expected)
;;
let run_test () =
test
~expected:(Some (V_Lit 42))
(app2 (app2 church_add (church 40) (church 2)) succ (lit 0));
test
~expected:(Some (V_Lit 42))
(app2
(app2 church_mul (church 3) (app2 church_add (church 9) (church 5)))
succ
(lit 0));
test
~expected:None
(app2 (lam @@ fun x -> lam @@ fun y -> app x y) (lit 1) (lit 2))
;;
end
let () = Lambda.run_test ()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment