extraction of small range types and for loops

- range types that fit in 31-bit signed integers are mapped to OCaml's type int
  (both literals and for loops)
- a for loop on a type that is mapped to 'int' in the OCaml driver is
  translated to an OCaml for loop

src/mlw/compile.ml

@@ -71,6 +71,7 @@ module ML = struct
         List.iter (fun (_, tyl) -> List.iter (iter_deps_ty f) tyl) constrl
     | Drecord pjl -> List.iter (fun (_, _, ty) -> iter_deps_ty f ty) pjl
     | Dalias ty -> iter_deps_ty f ty
+    | Drange _ | Dfloat _ -> ()
 
   let iter_deps_its_defn f its_d =
     Opt.iter (iter_deps_typedef f) its_d.its_def
@@ -340,8 +341,6 @@ module Translate = struct
           Mltree.Tapp (its.its_ts.ts_name, args) in
     loop t
 
-  let ty_int = mlty_of_ity MaskVisible ity_int
-
   let pvty pv =
     if pv.pv_ghost then ML.mk_var (pv_name pv) ML.tunit true
     else let (vs, vs_ty) = vsty pv.pv_vs in
@@ -633,11 +632,15 @@ module Translate = struct
       | Alias t, _, _ ->
           ML.mk_its_defn id args is_private (* FIXME ? is this a good mask ? *)
             (Some (Mltree.Dalias (mlty_of_ity MaskVisible t)))
-      | Range _, [], [] ->
+      | Range r, [], [] ->
+          assert (args = []); (* a range type is not polymorphic *)
+          ML.mk_its_defn id [] is_private (Some (Mltree.Drange r))
+      | Float ff, [], [] ->
           assert (args = []); (* a range type is not polymorphic *)
-          ML.mk_its_defn id [] is_private (Some (Mltree.Dalias ty_int))
-      | Range _, _, _ -> assert false (* a range type has no field/constr. *)
-      | Float _, _, _ -> assert false (* TODO *)
+          ML.mk_its_defn id [] is_private (Some (Mltree.Dfloat ff))
+      | (Range _ | Float _), _, _ ->
+          assert false (* cannot have constructors or fields *)
+
     end
 
   (* exception ExtractionVal of rsymbol *)

src/mlw/mltree.ml

@@ -88,6 +88,8 @@ type typedef =
   | Ddata     of (ident * ty list) list
   | Drecord   of (is_mutable * ident * ty) list
   | Dalias    of ty
+  | Drange    of Number.int_range
+  | Dfloat    of Number.float_format
 
 type its_defn = {
   its_name    : ident;

src/mlw/ocaml_printer.ml

@@ -283,6 +283,24 @@ module Print = struct
         Loc.errorm ?loc "Function %a cannot be extracted" Expr.print_rs rs
     | _ -> ()
 
+  (* a range type that fits in 31-bit sign integers is mapped to type "int" *)
+  let min_int31 = BigInt.of_string "-1073741824"
+  let max_int31 = BigInt.of_string "1073741823"
+  let is_small_range r =
+    BigInt.le min_int31 r.Number.ir_lower &&
+    BigInt.le r.Number.ir_upper max_int31
+  let is_small_range_ity ity = match ity.ity_node with
+    | Ityapp ({ Ity.its_def = Range r }, _, _) -> is_small_range r
+    | _ -> false
+  let is_small_range_type = function
+    | I ity -> is_small_range_ity ity | C _ -> false
+  let driver_type_int info ity = match ity.ity_node with
+    | Ityapp ({ its_ts = ts }, _, _) ->
+        query_syntax info.info_syn ts.ts_name = Some "int"
+    | _ -> false
+  let is_mapped_to_int info ity =
+    is_small_range_ity ity || driver_type_int info ity
+
   let print_constant fmt e = begin match e.e_node with
     | Econst c ->
         let s = BigInt.to_string (Number.compute_int_constant c) in
@@ -425,6 +443,7 @@ module Print = struct
           | _ -> assert false in
         (match query_syntax info.info_literal id with
          | Some s -> syntax_arguments s print_constant fmt [e]
+         | None when is_small_range_type e.e_ity -> pp_print_string fmt n
          | None   -> fprintf fmt (protect_on paren "Z.of_string \"%s\"") n)
     | Evar pvs ->
         (print_lident info) fmt (pv_name pvs)
@@ -498,24 +517,24 @@ module Print = struct
     | Eraise (xs, e_opt) ->
         print_raise ~paren info xs fmt e_opt
     | Efor (pv1, pv2, dir, pv3, e) ->
-        let for_id  = id_register (id_fresh "for_loop_to") in
-        let cmp, op = match dir with
-          | To -> "Z.leq", "Z.succ"
-          | DownTo -> "Z.geq", "Z.pred" in
-        fprintf fmt (protect_on paren
-               "@[<hov 2>let rec %a %a =@ if %s %a %a then \
-                begin@ %a; %a (%s %a) end@ in@ %a %a@]")
+        if is_mapped_to_int info pv1.pv_ity then
+          fprintf fmt "@[<hov 2>for %a = %a %a %a do@ @[%a@]@ done@]"
+            (print_lident info) (pv_name pv1) (print_lident info) (pv_name pv2)
+            print_for_direction dir (print_lident info) (pv_name pv3)
+            (print_expr info) e
+        else
+          let for_id  = id_register (id_fresh "for_loop_to") in
+          let cmp, op = match dir with
+            | To     -> "Z.leq", "Z.succ"
+            | DownTo -> "Z.geq", "Z.pred" in
+          fprintf fmt (protect_on paren
+                         "@[<hov 2>let rec %a %a =@ if %s %a %a then \
+                          begin@ %a; %a (%s %a) end@ in@ %a %a@]")
           (* let rec *) (print_lident info) for_id (print_pv info) pv1
           (* if      *)  cmp (print_pv info) pv1 (print_pv info) pv3
           (* then    *) (print_expr info) e (print_lident info) for_id
                         op (print_pv info) pv1
           (* in      *) (print_lident info) for_id (print_pv info) pv2
-        (*
-        fprintf fmt "@[<hov 2>for %a = %a %a %a do@ @[%a@]@ done@]"
-          (print_lident info) (pv_name pv1) (print_lident info) (pv_name pv2)
-          print_for_direction direction (print_lident info) (pv_name pv3)
-          (print_expr info) e
-        *)
     | Etry (e, bl) ->
         fprintf fmt
           "@[<hv>@[<hov 2>begin@ try@ %a@] with@]@\n@[<hov>%a@]@\nend"
@@ -582,6 +601,12 @@ module Print = struct
             (print_list newline print_field) fl
       | Some (Dalias ty) ->
           fprintf fmt " =@ %a" (print_ty ~paren:false info) ty
+      | Some (Drange r) when is_small_range r ->
+          fprintf fmt " =@ int"
+      | Some (Drange _) ->
+          fprintf fmt " =@ Z.t"
+      | Some (Dfloat _) ->
+          assert false (*TODO*)
     in
     let labels = its.its_name.id_label in
     if not (is_ocaml_remove ~labels) then

tests/test_extraction.mlw

@@ -213,12 +213,10 @@ module TestExtraction
 
   (** for loops *)
 
-(*
   let for_loop1 () =
-    let r = ref (0:int63) in
+    let r = ref (0: int63) in
     for i = 0 to 10 do r := !r + i done;
     !r
-*)
 
   (** optional and named arguments *)
 
@@ -300,6 +298,22 @@ module TestExtraction
 
   let test_foo () = foo 42
 
+  (* range types *)
+
+  type small_range = <range 1 10>
+
+  let test_small_range () =
+    for x = 1: small_range to 3 do
+      ()
+    done
+
+  type big_range = <range 1 0x4000_0000>
+
+  let test_big_range () =
+    for x = 1: big_range to 3 do
+      ()
+    done
+
 end
 
 (*