Partial Evaluation for Lambda Calculus

这是 [1] 第 8 章 Partial Evaluation for Lambda Calculus 的笔记.

在 Partial Evaluation for Functional Language 和 Partial Evaluation For Flow Chart Langauge 中, partial evaluation 所 eval 的东西很直观, 就是一个具体的像 int, bool 这样具体的值, 没有考虑高阶函数.

但是对于有高阶函数的语言, 情况变得复杂, 因为一个表达式的求值结果可能是一个函数, 那么考虑一个简单的场景, 返回一个常量的函数, 应该标记为是 Static 还是 Dynamic ? 比如: (lambda (x) 1)

如果标记为 S, 那这个函数在 residual program 中对应什么? 似乎也只能是 (lambda (x) 1)
如果标记为 D, 为什么一个这么简单的函数会返回一个常量的函数需要标记为 D? 是不是对于 lambda 表达式 partial evaluation 都无能为力?

我会有这样的疑惑主要有两个原因:

之前提到的 S 和 D 这样简单的 annotation 的标记对于简单的语言是足够的, 但是对于存在高阶函数的语言, 需要更丰富的 binding time annotation 才能描述"编译期函数"和"运行时函数";
标记为 D 和 S 的表达式都不一定会出现在最终的 residual program 中, 在有高阶函数的语言, residual program 长什么样主要看 partial evaluation 的结果, 之前的 “Dynamic 表达式作为程序骨架, Static 表达式求值后嵌入程序骨架” 的基本直觉失效了.

本文将介绍 lambda calculus 的 partial evaluation. 首先将定义一个简单的 lambda calculus, 然后介绍它的 binding time annotation 和 annotated version, 最后再分别展示它的 Binding Time Analysis(由 lambda calculus 得到 Annotated Program)和具体的 staging(由 annotated program 得到 residual program)算法.

本文所有代码都将用 OCaml 演示, 代码仓库在 https://github.com/butterunderflow/lambda_pe , 可以在这里 https://butter-xz.com/lambda_pe/ 在线体验 Lambda Calculus 的 Partial Evaluation;
有时候 Lambda Calulus 会被简写成 LC, Partial Evaluation 会被简写为 PE, Binding Time Analysis 会被简写为 BTA.

Lambda-calculus 的语法

labmda 表达式的语法定义如下:

(* expr1.ml, level 1 expression*)
type expr =
| EInt of int
| EVar of string
| ELam of string * expr
| ELet of string * expr * expr
| EApp of expr * expr

Binding Time Annotation(for Lambda Calculus)

有 partial evaluation 的语言有个 two-level type system, two-level 的意思是, 一个 level 对应编程语言原来的类型系统, 另一个就对应这个 annotation.

在之前章节里关于 Binding Time Annotation 的一个比较直观的理解是: 标记为 S 会被 eval 成值, D 会在 residual program 中被保留成代码.

这还是比较好理解的, 但是 D -> S 呢? D -> S -> S … 呢? 它们代表什么?

Binding time annotation 和一般的 type annotation 很类似, 只不过他不是描述 运行时(run time) 类型信息的, 而是描述 partial evaluation time 的类型信息的, 这个类型信息只做了 code 和 value 的区分, 具体的 code 和 value 又有它们自己在 first-level 的类型.

这个标记的所描述的是 pe-time 过程中的类型信息, 标记为 D(ynamic) 的表达式在 pe 过程中会被求值为 code, 标记为 S(tatic) 的会被求值为 value, 标记为 D -> S 的表达式在 specialize 过程中会转换为由 code 到 value 函数, 但是它们并不代表 residual program 中的最终形态. 也就是说, 标记为 S 和 D 的表达式都不一定会出现在 residual program 中.

Annotated Program(for Lambda Calculus)

Annotated Program 对应一个语言的 Two-level syntax

Figure 2: Lambda calculus 的 two-level syntax

带有 S 标记的表示这个表达式在 pe 过程中会 eval 到 value(可能是一般的值或函数值), 带有 D 标记的则会 eval 到 code.

note: variable 没有标记, variable 的 pe 结果是从当前的 environment 中 lookup 的结果.

(* expr2.ml, level 2 expression *)
type expr =
  (* use variable lookup as specialize result *)
  | Var of string
  (* specialize to a value(a pe time int or function value) *)
  | SConst of constant
  | SLam of string * expr
  | SLet of string * expr * expr
  | SApp of expr * expr
  (* specialize to a pe-time code expression *)
  | DLam of string * expr
  | DLet of string * expr * expr
  | DApp of expr * expr
  | DLift of expr
[@@deriving sexp]

Binding Time Analysis

LC 的 Binding Time Analysis 就是把 LC 编译到 2LC 的过程, 我们可以这样定义 BTA 的接口:

module type BTA_Sig = sig
  val analysis : Expr1.expr -> Expr2.expr
end

Naive BTA

BTA 是一个比较泛的概念, 并没有唯一正确的做法, 比如说什么都不干, 把所有东西都当作 Dynamic 也是一种可行的 BTA:

module NaiveBTA : BTA_Sig = struct
  (* blindly push every thing to runtime *)
  let analysis (e : E1.expr) : E2.expr =
    let rec go e =
      match e with
      | E1.EConst c -> E2.SConst c
      | E1.EVar x -> E2.Var x
      | E1.ELam (x, e0) -> E2.DLam (x, go e0)
      | E1.ELet (x, e0, e1) -> E2.DLet (x, go e0, go e1)
      | E1.EApp (e0, e1) -> E2.DApp (go e0, go e1)
    in
    go e
end

Naive BTA'

另一个极端是把所有东西都当作 static, 这么做当然也是可行的(只是没什么意义):

(* blindly stage every thing to compile time *)
let analysis (e : E1.expr) : E2.expr =
  let rec go e =
    match e with
    | E1.EConst c -> E2.SConst c
    | E1.EVar x -> E2.Var x
    | E1.ELam (x, e0) -> E2.SLam (x, go e0)
    | E1.ELet (x, e0, e1) -> E2.SLet (x, go e0, go e1)
    | E1.EApp (e0, e1) -> E2.SApp (go e0, go e1)
    | E1.EAnn (e0, _) -> go e0
    | E1.EOp (op, es) -> E2.SOp (op, List.map go es)
  in
  go e

BTA by type check

前面两个 NaiveBTA 是"对" 的吗? 当然是对的, 但是这其实没有意义, 因为我们没办法对 annotated program 做 staging (staging nothing 和 staging everything).

前面提到了, BTA 的结果是不唯一的, 既然不唯一, 那要如何选择一个表达式的 annotation 呢?

比如说下面的代码, f 既 apply 到了 1(a static)上, 又 apply 到了 y(a dynamic)上, 如何确定 f 的 annotation? 很多种可行解.

let f = fun x -> x in  (* f: D; D->D; S->S; (S->S) -> (S->S) ... *)
... f 1 ...
... f y ... (* y is dynamic *)

如果我们只看 f 1 , 那么 f 只能是 S -> 'a, D , D -> 'a (1可以被lift为Dynamic); 如果只看 f y, 那么 f 只能是 D -> 'a, D .

书上介绍了一种约束收集+求解的方法, 定义了一个 annotation 上的偏序关系, 然后通过求解收集到的约束得到最小的解.

这个方法太复杂了, 一堆偏序关系头都要晕了, 如果有感兴趣的话可以看 [1] 的 8.7 BTA by solving constraints .

这里选择另一种简单点的方法作为实现. 首先, 为了让 bta 变得简单, 我们给 LC1 加上一个 binding time annotation 注解的语法, 让我们可以手动的为表达式添加 annotation.

and expr =
  | EConst of constant
  | EVar of string
  | ELam of string * expr
  | ELet of string * expr * expr
  | EApp of expr * expr
  | EAnn of expr * Ann.t (* Binding time annotation hint *)
  | EOp of op * expr list

然后用 ML 类型推导的方式(不需要偏序关系, 仅仅需要等价关系)推导一个表达式的 Binding Time Annotation, 根据这个 Annotation 来得到 level 2 expression.

这里我们借用 Local Type Inference 中 bidirectional type checking 的思路, 把 bta 分为两种 mode:

check mode: 已知一个表达式的 annotation a, 验证这个表达式的 annotation 是否等于 a;
infer mode: 对表达式的 annotation 一无所知, 需要推导这个表达式的 annotation.

bidirectional type checking 可以把已经推导出的类型信息传递到相邻的语法树节点, 使得我们可以仅仅只写少量的 binding time annotation 注解就可以推导所有节点的 annotation.

具体的代码实现如下:

let rec infer (e : E1.expr) (env : ann_env) : E2.expr * ann =
  match e with
  | E1.EConst c -> (E2.SConst c, S)
  | E1.EVar x -> (E2.Var x, get x env)
  | E1.ELam (x, e0) -> failwith "can't infer a lambda"
  | E1.ELet (x, e0, e1) -> (
      let e0', a0' = infer e0 env in
      match a0' with
      | D -> (E2.DLet (x, e0', check e1 Ann.D ((x, a0') :: env)), D)
      | _ ->
          let e1', a1' = infer e1 ((x, a0') :: env) in
          (E2.SLet (x, e0', e1'), a1'))
  | E1.EApp (e0, e1) -> (
      let e0', a0' = infer e0 env in
      match a0' with
      | S -> failwith "error"
      | D ->
          let e1' = check e1 Ann.D env in
          (E2.DApp (e0', e1'), Ann.D)
      | Func (arg_ann, ret_ann) ->
          let e1' = check e1 arg_ann env in
          (E2.SApp (e0', e1'), ret_ann))
  | E1.EAnn (e0, a0) -> (check e0 a0 env, a0)
  | E1.EOp (op, es) -> (
      match (op, es) with
      | OAdd, [ e0; e1 ]
      | OMinus, [ e0; e1 ]
      | OAnd, [ e0; e1 ] ->
          let e0', a0' = infer e0 env in
          let e1' = check e1 a0' env in
          (E2.SOp (op, [ e0'; e1' ]), a0')
      | ONot, [ e0 ] ->
          let e0', a0' = infer e0 env in
          (E2.SOp (op, [ e0' ]), D)
      | _ -> failwith "neverreach")

and check (e : E1.expr) (a : ann) (env : ann_env) : E2.expr =
  match e with
  | E1.ELam (x, e0) -> check_lambda x e0 a env
  | _ ->
      let e', a' = infer e env in
      if a' = a then e'
      else if a' = S && a = D then E2.DLift e'
      else failwith "error"

and check_lambda x e a env =
  match a with
  | S -> failwith "error lambda annotation"
  | D ->
      let e' = check e D ((x, D) :: env) in
      E2.DLam (x, e')
  | Func (arg_ann, ret_ann) ->
      let e' = check e ret_ann ((x, arg_ann) :: env) in
      E2.SLam (x, e')

let analysis (e : E1.expr) : E2.expr = infer e empty_env |> fst

Staging!

万事俱备, 只欠 staging.

LC 的 two-level syntax 类似于之前在 partial evaluation for Scheme0 提到的 Binding Time Annotation , 但是在值域上有所差异:

Scheme0 的值域只能是 Constant 所在的值域(称之为 Const);
lambda calculus 的值域可以是高阶函数: 高阶函数的值域(2FuncVal)是很丰富的: 不仅可以是 \(Const \rightarrow Cosnt\) 的函数; 可以是 \(2FuncVal \rightarrow 2FuncVal\); 还可以是 \(Code\) … .

在做完 BTA 后, staging 的实现就很简单了, LC2 的 staging 其实也可以看作 two-level lambda calculus(2LC)的 interpreter! 只是在值域上有所差异, 这个 2LC 的值域在 LC 的基础上加上了 Code, 对于 2LC 表达式的求值结果可能是:

求值到 Const 或 FuncVal, 代表这个表达式已经求值完了;
还可能求值到一个 Code, 这个 Code 等之后再去算.

dlambda 的 body 也必须求值到 code, 然后再通过 build-lambda 创建一个动态求知 lambda 表达式. 但是这里为什么要用 newname? 书上说是为了避免 confusion, 我理解这个完全是为了 residual program 可读性和 specialize 算法的可维护性考虑的, 就算不重命名也不会导致 bug.

staging的代码实现如下:

let rec eval (e : expr) (env : env) : value =
match e with
| Var x -> List.assoc x env
| DLam (x, e) ->
    let new_var = gen_var ~hint:x in
    VCode (E1.ELam (x, eval e ((x, VCode new_var) :: env) |> get_code))
| DLet (x, e0, e1) ->
    let updated_env = (x, VCode (E1.EVar x)) :: env in
    VCode
      (match eval e0 env with
      | VCode code -> E1.ELet (x, code, eval e1 updated_env |> get_code)
      | VConst c ->
          E1.ELet (x, E1.EConst c, eval e1 updated_env |> get_code)
      | VFun f ->
          E1.ELet
            ( x,
              f (VCode (E1.EVar "_x")) |> get_code,
              eval e1 updated_env |> get_code ))
| DApp (e0, e1) ->
    VCode (E1.EApp (eval e0 env |> get_code, eval e1 env |> get_code))
| DLift e ->
    let v = get_int (eval e env) in
    VCode (E1.EConst v)
| SConst c -> VConst c
| SLam (x, e) -> VFun (fun v -> eval e ((x, v) :: env))
| SLet (x, e0, e1) ->
    let bind_value = eval e0 env in
    eval e1 ((x, bind_value) :: env)
| SApp (e0, e1) ->
    let func = get_func (eval e0 env) in
    eval e1 env |> func

thinking

Offline PE 中 annotated language(比如 2LC) 和 typeset language 的是否会存在某种同构?

References

[1]

N. D. Jones, C. K. Gomard, and P. Sestoft, Partial Evaluation and Automatic Program Generation. USA: Prentice-Hall, Inc., 1993.

Lambda-calculus 的语法#

Binding Time Annotation(for Lambda Calculus)#

Annotated Program(for Lambda Calculus)#

Binding Time Analysis#

Naive BTA#

Naive BTA'#

BTA by type check#

Staging!#

thinking#

References#