Adds relational operators.

1 files changed, 40 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