variance.ml

type variance = Bi | Co | Cn | Inv

type ty =
  | Unit
  | Tvar of int
  | Arrow of ty * ty
  | TyApp of int * ty array

type injectivity = bool

type tparam = variance * injectivity

type decl = tparam array * ty

let compute_variance =
  let compose = function
    | Bi, x | x, Bi -> x
    | Inv, _ | _, Inv -> Inv
    | Co, Co | Cn, Cn -> Co
    | _ -> Cn
  in
  let inv v = compose (v, Cn) in
  let union = function
    | Bi, x | x, Bi -> x
    | Co, Co -> Co
    | Cn, Cn -> Cn
    | _ -> Inv
  in
  let rec on_type env params v inj = function
    | Unit -> ()
    | Tvar i ->
      let (v', inj') = params.(i) in
      params.(i) <- (union (v, v'), inj || inj')
    | Arrow (t0, t1) ->
      on_type env params (inv v) inj t0;
      on_type env params v inj t1
    | TyApp (i, ts) ->
      let (decl_params, _) = env.(i) in
      Array.iter2 (fun ty (v', inj') ->
        on_type env params (compose (v ,v')) inj' ty
      ) ts decl_params
  in
  let on_decl env (params, ty) =
    let params = Array.copy params in
    on_type env params Co true ty;
    (params, ty)
  in
  let rec fixpoint env =
    let env' = Array.map (on_decl env) env in
    if env = env' then env else fixpoint env'
  in
  fixpoint

let show_variance (v, inj) =
  match v, inj with
  | _, false -> "Bi"
  | Bi, _ -> "Bi"
  | Co, _ -> "Co"
  | Cn, _ -> "Cn"
  | Inv, _ -> "Inv"

let pp_variance fmt v =
  Format.pp_print_string fmt (show_variance v)

let pp_decl fmt (params, _) =
  Format.fprintf fmt "[%a]" (Format.pp_print_array pp_variance) params

let pp_env fmt env =
  Format.fprintf fmt "@[%a@]" (Format.pp_print_array pp_decl) env

let run_test test =
  let env = Array.map (fun (n, ty) ->
    let params = Array.make n (Bi, false) in
    (params, ty)
  ) test in
  let env = compute_variance env in
  Format.printf "%a\n" pp_env env

(* T. T -> () *)
let test1 = [|(1, Arrow (Tvar 0, Unit))|]

(* T. () -> T *)
let test2 = [|(1, Arrow (Unit, Tvar 0))|]

(* T. (T -> ()) -> () *)
let test3 = [|(1, Arrow (Arrow (Tvar 0, Unit), Unit))|]

(* T. () *)
let test4 = [|(1, Unit)|]

(* A[T] = B[T] -> ()
   B[U] = A[U] -> () *)
let test5 =
  [|(1, Arrow (TyApp (1, [|Tvar 0|]), Unit));
    (1, Arrow (TyApp (0, [|Tvar 0|]), Unit))|]

(* A[T] = () -> B[T]
   B[U] = () -> U    *)
let test6 =
  [|(1, Arrow (Unit, TyApp (1, [|Tvar 0|])));
    (1, Arrow (Unit, Tvar 0))|]

let () =
  run_test test1;
  run_test test2;
  run_test test3;
  run_test test4;
  run_test test5;
  run_test test6;
  ()