A lifter for the toy language

toy.ml

open Core_kernel
open Bap_core_theory
open Bap.Std
open KB.Syntax
include Self()

module Word = struct
  include Bitvec_order
  include Bitvec_sexp.Functions
end

type ident = string [@@deriving compare,sexp]
type word = Word.t [@@deriving compare,sexp]
type sort = Word | Bool | Unit [@@deriving compare,sexp]
type ast =
  | Var of sort * ident
  | Int of word
  | Set of sort * ident * ast
  | App of sort * ident * ast
  | Rep of ast * ast
  | Seq of ast * ast
  | Ite of ast * ast * ast
  | Get of ast
  | Put of ast * ast
  | Nop
[@@deriving compare,sexp]

type def =
  | Proc of ident * ast
  | Glob of sort * ident * ast
  | Exec of ast
[@@deriving compare,sexp]

type prog = def list
[@@deriving compare,sexp]

let word = 8
let bool = 1

let pp_ast ppf x = Sexp.pp_hum ppf (sexp_of_ast x)

module Language = struct
  open Theory.Parser

  let bitv : type t r. (t,ast,r) bitv_parser =
    fun (module S) -> function
      | Var (Word,x) -> S.var x word
      | Int x -> S.int x word
      | Ite (c,x,y) -> S.ite c x y
      | Get x -> S.load Nop x
      | App (Word, "+", Seq (x,y)) -> S.add x y
      | App (Word, "-", Seq (x,y)) -> S.sub x y
      | App (Word, "*", Seq (x,y)) -> S.mul x y
      | App (Word, "/", Seq (x,y)) -> S.div x y
      | App (Word, "&", Seq (x,y)) -> S.logand x y
      | App (Word, "|", Seq (x,y)) -> S.logor x y
      | x ->
        error "bad word: %a"  pp_ast x;
        S.error

  let bool : type t r. (t,ast,r) bool_parser =
    fun (module S) -> function
      | Var (Bool,x) -> S.var x
      | Int x -> S.int x
      | Ite (c,x,y) -> S.ite c x y
      | App (Bool, "=", Seq (x,y)) -> S.eq x y
      | App (Bool, "<=", Seq (x,y)) -> S.le x y
      | App (Bool, "<", Seq (x,y)) -> S.lt x y
      | App (Bool, "<>", Seq (x,y)) -> S.neq x y
      | App (Bool, "&", Seq (x,y)) -> S.logand x y
      | App (Bool, "|", Seq (x,y)) -> S.logor x y
      | App (Bool, "not", x) -> S.not x
      | x ->
        error "bad bool: %a" pp_ast x;
        S.error

  let stmt : type t r. (t,ast,r,ast) stmt_parser =
    fun (module S) -> function
      | Set (Bool,v,x) -> S.set_bit v x
      | Set (Word,v,x) -> S.set_reg v word x
      | Put _ as m -> S.set_mem "mem" word word m
      | Rep (c,x) -> S.while_ c [x]
      | Ite (c,x,y) -> S.if_ c [x] [y]
      | Seq (x,y) -> S.seq [x;y]
      | App (Unit,name,Nop) -> S.call name
      | Nop -> S.seq []
      | x ->
        error "bad statement: %a" pp_ast x;
        S.error

  let mem : type t. (t,ast) mem_parser =
    fun (module S) -> function
      | Put (x,y) -> S.store Nop x y
      | _ -> S.error

  let rmode _ _ = assert false
  let float _ _ = assert false

  let t = {bitv; mem; stmt; bool; float; rmode}
end

let add_blks sub blks =
  List.fold blks ~init:sub ~f:(Term.append blk_t)

let build_program lift defs =
  let prog = Program.create () in
  Theory.Label.for_name "init" >>= fun tid ->
  let init = Sub.create ~tid ~name:"init" () in
  KB.List.fold defs ~init:(prog,init) ~f:(fun (prog,init) def ->
      match def with
      | Exec x ->
        lift x >>| fun blks ->
        (prog,add_blks init blks)
      | Glob (s,v,x) ->
        lift (Set (s,v,x)) >>| fun blks ->
        (prog,add_blks init blks)
      | Proc (name,body) ->
        Theory.Label.for_name name >>= fun tid ->
        let sub = Sub.create ~tid ~name () in
        lift body >>| fun blks ->
        let sub = add_blks sub blks in
        (Term.append sub_t prog sub,init))

let load filename : prog =
  Sexp.load_rev_sexps filename |>
  List.rev_map ~f:def_of_sexp


let lifter = KB.Class.declare "toy-language-lifter" ()
let program_t = KB.Domain.optional "program_t" ~equal:Program.equal
let program = KB.Class.property lifter "program" program_t
let source_t = KB.Domain.optional "source_t" ~equal:(fun x y ->
    compare_ast x y = 0)

let source = KB.Class.property Theory.Program.cls "source" source_t

let provide_semantics () =
  KB.promise Theory.Program.Semantics.slot @@ fun insn ->
  Theory.(instance>=>require) () >>= fun (module Core) ->
  let module Lifter = Theory.Parser.Make(Core) in
  KB.collect source insn >>= function
  | None -> KB.return Insn.empty
  | Some src -> Lifter.run Language.t [src]

let lift prog =
  let lift ast =
    KB.Object.create Theory.Program.cls >>= fun insn ->
    KB.provide source insn (Some ast) >>= fun () ->
    KB.collect Theory.Program.Semantics.slot insn >>| fun sema ->
    KB.Value.get Term.slot sema in
  build_program lift prog >>| fun (prog,init) ->
  let init = Term.set_attr init Sub.entry_point () in
  Term.prepend sub_t prog init

let run_lifter prog =
  let obj =
    KB.Object.create lifter >>= fun lifter ->
    lift prog >>= fun prog ->
    KB.provide program lifter (Some prog) >>| fun () ->
    lifter in
  match KB.run lifter obj (Toplevel.current ()) with
  | Error conflict ->
    error "ill-formed program: %a" KB.Conflict.pp conflict;
    invalid_arg "ill-formed program"
  | Ok (value,state) ->
    Toplevel.set state;
    match KB.Value.get program value with
    | None -> assert false
    | Some prog -> prog

let () = Project.Input.register_loader "toy" @@ fun filename ->
  let empty = Memmap.empty in
  let source = load filename in
  provide_semantics ();
  let prog = run_lifter source in
  Project.Input.create `unknown filename ~code:empty ~data:empty
    ~finish:(fun proj -> Project.with_program proj prog)