type int_size = IPow of char | ISize

type ty_lit =
  | Type
  | Bool
  | Char
  | String
  | TIntSize
  | UnsInt of int_size
  | SigInt of int_size

type lit =
  | LTy of ty_lit
  | LIntSize of int_size
  | LNatural of int

let lty x = LTy x

let isz_of_str = function
  | "1"   -> IPow '\000' |> Option.some
  | "2"   -> IPow '\001' |> Option.some
  | "4"   -> IPow '\002' |> Option.some
  | "8"   -> IPow '\003' |> Option.some
  | "16"  -> IPow '\004' |> Option.some
  | "32"  -> IPow '\005' |> Option.some
  | "64"  -> IPow '\006' |> Option.some
  | "128" -> IPow '\007' |> Option.some
  | "256" -> IPow '\008' |> Option.some
  | "_size" -> ISize |> Option.some
  | _ -> Option.none

let isz_to_str = function
  | IPow x -> Int.shift_left 1 (Char.code x) |> Int.to_string
  | ISize -> "_size"

let isz_to_pow = function
  | IPow x -> x
  | ISize -> 16 (* should be a fitting default *)

let str_post loc s = String.sub s loc (String.length s - loc)

let of_string = function
  | "Type"    -> LTy Type     |> Option.some
  | "Bool"    -> LTy Bool     |> Option.some
  | "Char"    -> LTy Char     |> Option.some
  | "String"  -> LTy String   |> Option.some
  | "IntSize" -> LTy TIntSize |> Option.some

  | s ->  (if String.starts_with ~prefix:"ZI" s then
               str_post 2 s |> isz_of_str |> Option.map (fun x -> LIntSize x)
      else if String.starts_with ~prefix:"UI" s then
               str_post 2 s |> isz_of_str |> Option.map (fun x -> UnsInt x |> lty)
      else if String.starts_with ~prefix:"SI" s then
               str_post 2 s |> isz_of_str |> Option.map (fun x -> SigInt x |> lty)
      else if String.starts_with ~prefix:"-"  s then Option.none
      else match int_of_string_opt s with
        | Some x -> Some (LNatural x)
        | None -> None)

let to_string = function
  | LTy Type     -> "Type"
  | LTy Bool     -> "Bool"
  | LTy Char     -> "Char"
  | LTy String   -> "String"
  | LTy TIntSize -> "IntSize"
  | LIntSize sz     -> "ZI" ^ (isz_to_str sz)
  | LTy (UnsInt sz) -> "UI" ^ (isz_to_str sz)
  | LTy (SigInt sz) -> "SI" ^ (isz_to_str sz)
  | LNatural n      -> Int.to_string n