Adds relational operators.

3 files changed, 76 insertions, 0 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index 8e9c0c9..a4e2272 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -13,10 +13,21 @@ type 'a index = { index : int } [@@deriving show { with_path = false }]
 
 let get (arr : 'a Arraylist.t) (i : 'a index) : 'a = Arraylist.get arr i.index
 
+type relop = [ `LT | `LTE | `GT | `GTE | `EQ | `NE ]
+
+let string_of_relop : relop -> string = function
+  | `LT -> "<"
+  | `LTE -> "<="
+  | `GT -> ">"
+  | `GTE -> ">="
+  | `EQ -> "="
+  | `NE -> "<>"
+
 type 'a prim =
   | Add : (expr index * expr index) prim
   | Mul : (expr index * expr index) prim
   | Sub : (expr index * expr index) prim
+  | RelOp : (relop * expr index * expr index) prim
 
 and expr =
   | App of expr index * expr index
@@ -68,6 +79,30 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
         ( { pexp_desc = Pexp_ident { txt = Lident "*"; _ }; _ },
           [ (Nolabel, l); (Nolabel, r) ] ) ->
         add (Prim (Mul, (loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "<"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`LT, loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "<="; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`LTE, loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident ">"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`GT, loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident ">="; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`GTE, loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "="; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`EQ, loop l, loop r)))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "<>"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (RelOp, (`NE, loop l, loop r)))
     | Pexp_apply (f, xs) ->
         List.fold_left
           (fun f -> function
@@ -213,6 +248,11 @@ let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
           (Pexp_apply
              ( Wrap.expression (Pexp_ident (Wrap.ident "*")),
                [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
+    | Prim (RelOp, (op, l, r)) ->
+        Wrap.expression
+          (Pexp_apply
+             ( Wrap.expression (Pexp_ident (Wrap.ident (string_of_relop op))),
+               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
     | Var name -> Wrap.expression (Pexp_ident (Wrap.ident name))
   in
   loop
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 08a8288..685ef52 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -25,6 +25,9 @@ let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
   | Prim (Add, _) -> Format.fprintf fmt "  expr%d [label=\"+\"];\n" i.index
   | Prim (Sub, _) -> Format.fprintf fmt "  expr%d [label=\"-\"];\n" i.index
   | Prim (Mul, _) -> Format.fprintf fmt "  expr%d [label=\"*\"];\n" i.index
+  | Prim (RelOp, (op, _, _)) ->
+      Format.fprintf fmt "  expr%d [label=\"%s\"];\n" i.index
+        (string_of_relop op)
   | Var x -> Format.fprintf fmt "  expr%d [label=%S];\n" i.index x
 
 let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
@@ -64,6 +67,9 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
     | Prim (Mul, (l, r)) ->
         edge_to l;
         edge_to r
+    | Prim (RelOp, (_, l, r)) ->
+        edge_to l;
+        edge_to r
     | Var _ -> ()
     | Bool _ | Int _ | Nil -> ()
   in
diff --git a/discocaml/eval.ml b/discocaml/eval.ml
index f1e5a2b..cdcb727 100644
--- a/discocaml/eval.ml
+++ b/discocaml/eval.ml
@@ -30,6 +30,9 @@ let has_free_var (ast : expr ast) (i : expr index) (name : string) : bool =
     | Prim (Mul, (l, r)) ->
         loop l;
         loop r
+    | Prim (RelOp, (_, l, r)) ->
+        loop l;
+        loop r
     | Var x -> if String.equal name x then raise FoundIt
   in
   try
@@ -75,6 +78,11 @@ let check_redex (ast : expr ast) (i : expr index) : unit =
       match (get_subexpr ast l, get_subexpr ast r) with
       | Int _, Int _ -> raise (FoundRedex i)
       | _ -> ())
+  | Prim (RelOp, (_, l, r)) -> (
+      match (get_subexpr ast l, get_subexpr ast r) with
+      | Bool _, Bool _ -> raise (FoundRedex i)
+      | Int _, Int _ -> raise (FoundRedex i)
+      | _ -> ())
   | Var _ -> ()
   | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ()
 
@@ -89,6 +97,9 @@ let find_redex_cbv_in (ast : expr ast) : expr index -> unit =
     | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
         loop l;
         loop r
+    | Prim (RelOp, (_, l, r)) ->
+        loop l;
+        loop r
     | Var _ -> ()
     | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ());
     check_redex ast i
@@ -107,6 +118,9 @@ let find_redex_cbn_in (ast : expr ast) : expr index -> unit =
     | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
         loop l;
         loop r
+    | Prim (RelOp, (_, l, r)) ->
+        loop l;
+        loop r
     | Var _ -> ()
     | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ()
   in
@@ -151,6 +165,7 @@ let subst (ast : expr ast) (from : string) (to_ : expr index) :
     | Prim (Add, (l, r)) -> add (Prim (Add, (loop l, loop r)))
     | Prim (Sub, (l, r)) -> add (Prim (Sub, (loop l, loop r)))
     | Prim (Mul, (l, r)) -> add (Prim (Mul, (loop l, loop r)))
+    | Prim (RelOp, (op, l, r)) -> add (Prim (RelOp, (op, loop l, loop r)))
     | Var x -> if String.equal from x then to_ else i
   in
   loop
@@ -177,5 +192,20 @@ let reduce (ast : expr ast) (i : expr index) : expr ast =
     | Prim (Add, (l, r)) -> Int (must_int l + must_int r)
     | Prim (Sub, (l, r)) -> Int (must_int l - must_int r)
     | Prim (Mul, (l, r)) -> Int (must_int l * must_int r)
+    | Prim (RelOp, (op, l, r)) ->
+        let l, r =
+          match (get_subexpr ast l, get_subexpr ast r) with
+          | Bool l, Bool r -> (Obj.repr l, Obj.repr r)
+          | Int l, Int r -> (Obj.repr l, Obj.repr r)
+          | _ -> fail ()
+        in
+        Bool
+          (match op with
+          | `LT -> l < r
+          | `LTE -> l <= r
+          | `GT -> l > r
+          | `GTE -> l >= r
+          | `EQ -> l = r
+          | `NE -> l <> r)
     | _ -> fail ());
   ast