Commits

bergsoe committed 9c00d04

Micro optimizations.

Prefer the raw [next] method of BatEnum to the
BatEnum.get function that wraps its results.

  • Participants
  • Parent commits 72db8cb

Comments (0)

Files changed (1)

src/PaplSampler.ml

 
 let stop () = raise BatEnum.No_more_elements
 
+let next_fun s =
+  let obj = BatEnum.to_object s in
+    fun () -> obj#next
+
 let pi = BatFloat.pi
 
+let map2 f sa sb =
+  let next_a = next_fun sa in
+  let next_b = next_fun sb in
+  let next () =
+    let a = next_a () in
+    let b = next_b () in
+      f a b
+  in
+    BatEnum.from next
+
 let product2 sa sb = BatEnum.combine (sa, sb)
 
 let product3 sa sb sc =
+  let next_a = next_fun sa in
+  let next_b = next_fun sb in
+  let next_c = next_fun sc in
   let next () =
-    match BatEnum.get sa, BatEnum.get sb, BatEnum.get sc with
-        Some a, Some b, Some c -> (a, b, c)
-      | _ -> stop ()
+    let a = next_a () in
+    let b = next_b () in
+    let c = next_c () in
+      (a, b, c)
   in
     BatEnum.from next
 
 let intersperse sx total sy =
   let n = ref total in
   let cnt = ref 0 in
+  let next_x = next_fun sx in
   let rec next () =
     if !cnt == 0 then
-      match BatEnum.get sx with
-          None -> stop ()
-        | Some x ->
-            cnt := !n;
-            x
+      let x = next_x () in
+        cnt := !n;
+        x
     else begin
       cnt := !cnt - 1;
       match BatEnum.get sy with
   BatEnum.filter (PaplConstraint.accept constr) sampler
 
 let rec constrain_with_default sampler constr default =
+  let next_sample = next_fun sampler in
   let rec next () =
-    match BatEnum.get sampler with
-        None -> stop ()
-      | Some x ->
-          if c constr x then
-            match BatEnum.get default with
-                None -> next ()
-              | Some y -> y
-          else
-            x
+    let x = next_sample () in
+      if c constr x then
+        match BatEnum.get default with
+            None -> next ()
+          | Some y -> y
+      else
+        x
   in BatEnum.from next
 
 let rec round_robin samplers =
           lookup (right i) p'
         else
           snd (get i) in
-  let make () =
-    match BatEnum.get position_sampler with
-        None -> stop ()
-      | Some p -> lookup 0 p
-  in
+  let next_position = next_fun position_sampler in
+  let next () = lookup 0 (next_position ()) in
     if Array.length tree == 0 then
       invalid_arg "Empty enum of pairs given."
     else
-      BatEnum.from make
+      BatEnum.from next
 
 let distribute_by_weight ?rng pairs =
   distribute_by_weight_helper (PaplRandom.enum_float ?rng) pairs