#light

open StructuredFormat
open StructuredFormat.LayoutOps

let get_layout (env: #IEnvironment) (e: #IFormattable) = e.GetLayout(env)

type ast =
 | Val of int
 | Var of string
 | BinOp of binop * ast * ast
 | UnOp of unop * ast
with
   static member (+)(a, b)  = BinOp (Plus, a, b)
   static member (-)(a, b)  = BinOp (Minus, a, b)
   static member ( *)(a, b) = BinOp (Times, a, b)
   static member (/)(a, b)  = BinOp (Div, a, b)
   static member (~+)(a)    = UnOp (UPlus, a)
   static member (~-)(a)    = UnOp (UMinus, a)
   interface StructuredFormat.IFormattable with
      member x.GetLayout(env) =
         match x with
          | Val i -> objL (box i)
          | Var s -> objL (box s)
          | BinOp (Custom (s, _), t1, t2) ->
               wordL s ++ get_layout env t1 ++ get_layout env t2
          | BinOp (b, t1, t2) ->
               (get_layout env t1 $$ get_layout env b) -- get_layout env t2 |> bracketL
          | UnOp (b, t1) ->
               get_layout env b $$ get_layout env t1
   end
end

and binop =
 | Plus
 | Minus
 | Times
 | Div
 | Custom of string * (int -> int -> int)
with
   interface StructuredFormat.IFormattable with
      member x.GetLayout(env) =
         match x with
          | Plus -> wordL "+"
          | Minus -> wordL "-"
          | Times -> wordL "*"
          | Div -> wordL "/"
          | Custom (s, _) -> wordL s
   end
end

and unop = UPlus | UMinus
with
   interface StructuredFormat.IFormattable with
      member x.GetLayout(env) =
         match x with
          | UPlus -> leftL "+"
          | UMinus -> leftL "-"
   end
end

// Apply a function f on every node of the tree
let rec map f = function
  | Val _ | Var _ as a -> f a
  | BinOp (op, a, b)   -> f (BinOp (op, map f a, map f b))
  | UnOp (op, a)       -> f (UnOp (op, map f a))

// Simplify an AST
//   do the evaluation when possible
//   use arithmetic rules when there are unkown values
let simplify =
  let op_to_fct = function
    | Plus -> (+) | Minus -> (-) | Times -> ( * ) | Div -> (/)
    | Custom (_, a) -> a
  in map (function
    | BinOp (op, Val a, Val b) -> Val ((op_to_fct op) a b)
    | UnOp (UMinus, Val a) -> Val (-a)
    | BinOp (Times, Val 0, _) | BinOp (Times, _, Val 0) -> Val 0
    | BinOp (Times, Val 1, a) | BinOp (Times, a, Val 1)
    | BinOp (Minus, a, Val 0) | BinOp (Div, a, Val 1)
    | UnOp (UPlus, a)         | UnOp (_, UnOp (_, a))
    | a -> a)

let (!) a = Val a
let (!?) a = Var a
let min' a b = BinOp(Custom ("min", min), a, b)

// Example
let () =
   // expression: - (-x) * (3 + min(5, 0) * x + y)
   let ast = - (- !?"x") * (!3 + (min' !5 !0) * !?"x") + !?"y"
   // simplifies to: x * 3 + y
   in print_any (ast |> simplify)