Implement Float shift methods in C

2 files changed, 51 insertions, 29 deletions

diff --git a/mrblib/numeric.rb b/mrblib/numeric.rb
index 6e4c5027f..17155bfd6 100644
--- a/mrblib/numeric.rb
+++ b/mrblib/numeric.rb
@@ -160,35 +160,7 @@ end
 #
 # ISO 15.2.9
 class Float
-  include Integral
   # mruby special - since mruby integers may be upgraded to floats,
   # floats should be compatible to integers.
-  def >> other
-    n = self.to_i
-    other = other.to_i
-    if other < 0
-      n << -other
-    else
-      other.times { n /= 2 }
-      if n.abs < 1
-        if n >= 0
-          0
-        else
-          -1
-        end
-      else
-        n.to_i
-      end
-    end
-  end
-  def << other
-    n = self.to_i
-    other = other.to_i
-    if other < 0
-      n >> -other
-    else
-      other.times { n *= 2 }
-      n
-    end
-  end
+  include Integral
 end
diff --git a/src/numeric.c b/src/numeric.c
index cdf8c8afb..c86373318 100644
--- a/src/numeric.c
+++ b/src/numeric.c
@@ -15,6 +15,7 @@
 #include <mruby/string.h>
 
 #ifdef MRB_USE_FLOAT
+#define trunc(f) truncf(f)
 #define floor(f) floorf(f)
 #define ceil(f) ceilf(f)
 #define fmod(x,y) fmodf(x,y)
@@ -342,6 +343,53 @@ flo_xor(mrb_state *mrb, mrb_value x)
   return int64_value(mrb, v1 ^ v2);
 }
 
+static mrb_value
+flo_shift(mrb_state *mrb, mrb_value x, mrb_int width)
+{
+  mrb_float val;
+
+  if (width == 0) {
+    return x;
+  }
+  val = mrb_float(x);
+  if (width < 0) {
+    while (width++) {
+      val /= 2;
+    }
+    val = trunc(val);
+    if (val == 0 && mrb_float(x) < 0) {
+      return mrb_fixnum_value(-1);
+    }
+  }
+  else {
+    while (width--) {
+      val *= 2;
+    }
+  }
+  if (FIXABLE(val)) {
+    return mrb_fixnum_value(val);
+  }
+  return mrb_float_value(mrb, val);
+}
+
+static mrb_value
+flo_lshift(mrb_state *mrb, mrb_value x)
+{
+  mrb_int width;
+
+  mrb_get_args(mrb, "i", &width);
+  return flo_shift(mrb, x, -width);
+}
+
+static mrb_value
+flo_rshift(mrb_state *mrb, mrb_value x)
+{
+  mrb_int width;
+
+  mrb_get_args(mrb, "i", &width);
+  return flo_shift(mrb, x, width);
+}
+
 /* 15.2.8.3.18 */
 /*
  * call-seq:
@@ -1253,6 +1301,8 @@ mrb_init_numeric(mrb_state *mrb)
   mrb_define_method(mrb, fl,      "&",         flo_and,          MRB_ARGS_REQ(1));
   mrb_define_method(mrb, fl,      "|",         flo_or,           MRB_ARGS_REQ(1));
   mrb_define_method(mrb, fl,      "^",         flo_xor,          MRB_ARGS_REQ(1));
+  mrb_define_method(mrb, fl,      ">>",        flo_lshift,       MRB_ARGS_REQ(1));
+  mrb_define_method(mrb, fl,      "<<",        flo_rshift,       MRB_ARGS_REQ(1));
   mrb_define_method(mrb, fl,      "ceil",      flo_ceil,         MRB_ARGS_NONE()); /* 15.2.9.3.8  */
   mrb_define_method(mrb, fl,      "finite?",   flo_finite_p,     MRB_ARGS_NONE()); /* 15.2.9.3.9  */
   mrb_define_method(mrb, fl,      "floor",     flo_floor,        MRB_ARGS_NONE()); /* 15.2.9.3.10 */