open Core_kernel[@@warning "-D"]
open Bap.Std
open Regular.Std
open Mips_rtl
open Mips_utils

type 'a ec = bool -> 'a

type bitwidth = int

let bit = 1
let byte = 8
let halfword = 16
let word = 32
let doubleword = 64
let quadword = 128
let bitwidth x = x

let int_of_bitwidth = Fn.id

let width_of_size = Size.in_bits

let int_of_imm = function
  | Op.Reg _ | Op.Fmm _ -> mips_fail "imm operand expected"
  | Op.Imm x -> match Imm.to_int x with
    | Some x -> x
    | None -> mips_fail "failed to convert imm operand to int"

let imm signed op =
  let w = Word.of_int ~width:32 (int_of_imm op) in
  if signed then Exp.(signed @@ of_word w)
  else Exp.(unsigned @@ of_word w)

let signed f = f true
let unsigned f = f false

let apply_signess signed e =
  if signed then Exp.signed e
  else Exp.unsigned e

let var signed width =
  Exp.tmp (int_of_bitwidth width) |>
  apply_signess signed

let reg find signed op = match op with
  | Op.Imm _ | Op.Fmm _ -> mips_fail "reg operand expected"
  | Op.Reg x -> apply_signess signed (find x)

let const signed width value =
  let width = int_of_bitwidth width in
  let x = Word.of_int ~width value in
  apply_signess signed (Exp.of_word x)

let of_string signed s =
  let s = String.filter ~f:(fun c -> Char.(c <> '_')) s in
  let chop (prefix, len) =
    match String.chop_prefix ~prefix s with
    | None -> None
    | Some data -> Some (len,data) in
  let width,data =
    match List.find_map ~f:chop ["0x",4; "0o",3; "0b",1;] with
    | Some (len, data) -> String.length data * len, data
    | None -> Z.numbits (Z.of_string s), s in
  let suf = if signed then "s" else "u" in
  let w = Word.of_string (sprintf "%s:%d%s" s width suf) in
  apply_signess signed (Exp.of_word w)

let zero = Exp.of_word Word.b0
let one  = Exp.of_word Word.b1

let last e bits =
  let w = Exp.width e in
  Exp.extract (w - 1) (w - bits) e

let first e bits = Exp.extract (bits - 1) 0 e

let high w e = last e (int_of_bitwidth w)
let low w e = first e (int_of_bitwidth w)

let msb e =
  let h = Exp.width e - 1 in
  Exp.extract h h e

let lsb e = Exp.extract 0 0 e

let nth w e index =
  let step = int_of_bitwidth w in
  let hi = (index + 1) * step - 1 in
  let lo = index * step in
  Exp.extract hi lo e

let extract e hi lo = Exp.extract hi lo e

let when_ cond then_ = if_ cond then_ []
let ifnot cond else_ = if_ cond [] else_

type clause = [
  | `Case of (exp * rtl list)
  | `Default of rtl list
]

let case x y = `Case (x,y)
let default y = `Default y

let switch exp cases =
  let default = List.filter_map ~f:(function
      | `Default y -> Some y
      |  _ -> None) cases in
  let default = Option.value ~default:[] (List.hd default) in
  let cases = List.filter_map ~f:(function
      | `Case (x,y) -> Some (x,y)
      | _ -> None) cases in
  let cond x = Infix.(exp = x) in
  match cases with
  | [] -> mips_fail "empty switch"
  | (x, code) :: [] -> (if_ (cond x) code default;)
  | (x, code) :: cases ->
    let else_ =
      List.fold (List.rev cases) ~init:default ~f:(fun acc (x,code) ->
          [if_ (cond x) code acc;]) in
    (if_ (cond x) code else_)

let width e =
  let w = Exp.width e in
  Exp.of_word (Word.of_int ~width:w w)
