(* Naïve closure conversion (for single-argument functions, with
   runtime value environments represented as nested pairs). *)

(* The source language. *)

type source =
  | SVar of atom
  | SAbs of < atom * inner source >
  | SApp of source * source

(* The target language has closed, binary functions, and pairs. *)

type target =
  | TVar of atom
  | TAbs of < xve: atom * x: atom * inner t: target >
            where free(x) # free(xve)            (* the two formal arguments have distinct names *)
              and free(t) <= free(x) U free(xve) (* the function is closed *)
  | TApp of target * target * target
  | TUnit
  | TPair of target * target
  | TFst of target
  | TSnd of target
  | TLet of < atom * outer target * inner target >

(* Translation environments are ordered lists of atoms. *)

type env =
  | ENil
  | ECons of x: atom * e: env
             where free(x) # free(e)             (* no duplicate elements *)

(* Turning a set of atoms into a list of atoms. (Sets of atoms are a
   primitive type. [empty?] and [choose] are primitive operations.) *)

fun linearize accepts s produces xs where free(s) == free(xs) =
  if empty? (s) then
    ENil
  else
    let x, s = choose (s) in
    ECons (x, linearize (s))
  end

(* Constructing a list of an expression's free variables. We exploit
   the primitive operation [free], which produces a set of an
   expression's free variables, and turn that into a list. One could
   also do everything with lists, without using any primitive
   operations. *)

fun fv accepts e produces xs where free(xs) == free(e) =
  linearize (free (e))

(* Code production for environment lookup. [env] is the translation
   environment, which can be viewed as a mapping of variables to
   integers. [xve] is an expression that denotes the runtime value
   environment. [x] is the variable that we are looking up. *)

fun lookup
  accepts env, xve, x
    where free(x) <= free(env)
  produces t
    where free(t) <= free(xve)
=
  case env of
  | ENil ->
      absurd
  | ECons (y, tail) ->
      if x == y then
        (* project the first element of the runtime value environment *)
	TFst (xve)
      else
	(* skip one element of the runtime value environment *)
	lookup (tail, TSnd (xve), x)
      end
  end

(* Translating variable accesses. [env] and [xve] are as above. [arg]
   is the name of the current function's argument -- it is the only
   variable that can be directly accessed. [SVar x] is the source
   expression that we wish to translate. *)

fun convert_variable
  accepts env, xve, arg, x
    where free(x) <= free(env) U free(arg)
  produces t
    where free(t) <= if free(env) != empty then free(xve) end U (free(x) @ free(arg))
    (* The conditional means that [xve] can appear in the result only if
       [env] is non=empty. *)
    (* The intersection means that [arg] can appear in the result only if
       [x] is [arg]. *)
=
  if x == arg then
    TVar (arg)
  else
    lookup (env, xve, x)
  end

(* Code production for environment construction. [env] and [xve] are
   as above. [xs] is a list of variables that need to be accessed.
   The results of these individual accesses conceptually form a list,
   which we represent as a target term using nested pairs. *)

fun build
  accepts env, xve, arg, xs
    where free(xs) <= free(env) U free(arg)
  produces t
    where free(t) <= if free(env) != empty then free(xve) end U (free(xs) @ free(arg))
=
  case xs of
  | ENil ->
      TUnit
  | ECons (x, xs) ->
      TPair (convert_variable (env, xve, arg, x), build (env, xve, arg, xs))
  end

(* Conversion. [env] and [xve] are as above. [arg] is the name of the
   current function's argument -- it is the only variable that can be
   accessed directly. [e] is the expression that we wish to
   translate. *)

fun convert
  accepts env, xve, arg, e
    where free(xve) # free(arg)
      and free(e) <= free(env) U free(arg)
  produces t
    where free(t) <= if free(env) != empty then free(xve) end U (free(e) @ free(arg))
=
  case e of
  | SVar (x) ->

      convert_variable (env, xve, arg, x)

  | SAbs (z, body) ->

      (* Create a new translation environment by examining the
	 free variables of [e]. *)

      let env' = fv (e) in

      (* Translate each variable within [env'] under [env], thus
	 creating a tuple that will be used as the closure's runtime
	 value environment. *)

      let ve = build (env, xve, arg, env') in

      (* Generate a fresh variable to stand for the runtime value
	 environment inside the translated function. *)

      fresh xve in

      (* Translate the function body under the new translation
	 environment. *)

      let t = convert (env', TVar (xve), z, body) in

      (* Construct a pair of the target function and its value
	 environment. *)

      TPair (TAbs (xve, z, t), ve)

  | SApp (e1, e2) ->

      (* Evaluate the closure, extract the code pointer and the value
	 environment, perform the application of the code pointer to
	 the value environment and the actual argument. *)

      fresh closure in
      TLet (closure,
	    convert (env, xve, arg, e1),
	    TApp (TFst (TVar (closure)), TSnd (TVar (closure)), convert (env, xve, arg, e2)))

  end

(* Converting a closed expression. Since the expression is closed, we
   really don't need the [xve] and [arg] parameters. We generate two
   dummies, which we pass to [convert]. The difficulty is then to
   prove that these two dummies do not appear in [convert]'s
   result. This works because [convert]'s postcondition is precise
   enough. The postcondition states: (i) that [xve] can appear in the
   result only if [env] is non-empty, and (ii) that [arg] can appear
   in the result only if [arg] appears in [e]. Here, neither condition
   is met, so the result is closed. *)

fun cc accepts e where free(e) == empty produces t =
  fresh xve, arg in
  convert (ENil, TVar (xve), arg, e)