simplified shift operators

backend/cmm_helpers.ml

@@ -416,24 +416,35 @@ let lsr_int c1 c2 dbg =
   | Cop (Clsr, [c; Cconst_int (n1, _)], _), Cconst_int (n2, _)
     when n1 > 0 && n2 > 0 && n1 + n2 < size_int * 8 ->
     Cop (Clsr, [c; Cconst_int (n1 + n2, dbg)], dbg)
+  | Cop (Clsr, [c; Cconst_int (n, _)], _), c2 when 0 < n && n < arch_bits ->
+    (* prefer to keep the constant shift on the outside to enable further
+       optimizations. *)
+    Cop (Clsr, [Cop (Clsr, [c; c2], dbg); Cconst_int (n, dbg)], dbg)
   | c1, Cconst_int (n, _) when n > 0 ->
     Cop (Clsr, [ignore_low_bit_int c1; c2], dbg)
   | _ -> Cop (Clsr, [c1; c2], dbg)
 
 let lsr_const c n dbg = lsr_int c (Cconst_int (n, dbg)) dbg
 
 let asr_int c1 c2 dbg =
-  match c2 with
-  | Cconst_int (0, _) -> c1
-  | Cconst_int (n, _) when n > 0 -> (
+  match c1, c2 with
+  | c1, Cconst_int (0, _) -> c1
+  | Cop (Casr, [c; Cconst_int (n1, _)], _), Cconst_int (n2, _)
+    when n1 > 0 && n2 > 0 && n1 + n2 < size_int * 8 ->
+    Cop (Casr, [c; Cconst_int (n1 + n2, dbg)], dbg)
+  | c1, Cconst_int (n, _) when n > 0 -> (
     match ignore_low_bit_int c1 with
     (* some operations always return small enough integers that it is safe and
        correct to optimise [asr (lsl x 1) 1] into [x]. *)
     | Cop (Clsl, [c; Cconst_int (1, _)], _)
       when n = 1 && guaranteed_to_be_small_int c ->
       c
     | c1' -> Cop (Casr, [c1'; c2], dbg))
-  | _ -> Cop (Casr, [c1; c2], dbg)
+  | Cop (Casr, [c; Cconst_int (n, _)], _), c2 when 0 < n && n < arch_bits ->
+    (* prefer to keep the constant shift on the outside to enable further
+       optimizations. *)
+    Cop (Casr, [Cop (Casr, [c; c2], dbg); Cconst_int (n, dbg)], dbg)
+  | c1, c2 -> Cop (Casr, [c1; c2], dbg)
 
 let asr_const c n dbg = asr_int c (Cconst_int (n, dbg)) dbg
 
@@ -1321,15 +1332,19 @@ let rec low_bits ~bits x dbg =
     bits set to 0 *)
 let zero_extend ~bits e dbg =
   assert (0 < bits && bits <= arch_bits);
-  let mask = Nativeint.pred (Nativeint.shift_left 1n bits) in
-  let zero_extend_via_mask e =
-    Cop (Cand, [e; natint_const_untagged dbg mask], dbg)
+  let unused_bits = arch_bits - bits in
+  let zero_extend_via_shift e =
+    lsr_const (lsl_const e unused_bits dbg) unused_bits dbg
   in
   if bits = arch_bits
   then e
   else
     map_tail
       (function
+        | Cop (Clsr, [_; Cconst_int (n, _)], _)
+          when unused_bits <= n && n < arch_bits ->
+          (* already has zero in the high bits *)
+          e
         | Cop (Cload { memory_chunk; mutability; is_atomic }, args, dbg) as e
           -> (
           let load memory_chunk =
@@ -1340,8 +1355,8 @@ let zero_extend ~bits e dbg =
           | (Sixteen_signed | Sixteen_unsigned), 16 -> load Sixteen_unsigned
           | (Thirtytwo_signed | Thirtytwo_unsigned), 32 ->
             load Thirtytwo_unsigned
-          | _ -> zero_extend_via_mask e)
-        | e -> zero_extend_via_mask e)
+          | _ -> zero_extend_via_shift e)
+        | e -> zero_extend_via_shift e)
       (low_bits ~bits e dbg)
 
 let sign_extend ~bits e dbg =
@@ -4484,11 +4499,11 @@ module Numeric = struct
 
     let nativeint = Untagged Integer.nativeint
 
+    let[@inline] untagged = function
+      | Untagged t -> t
+      | Tagged t -> Tagged_integer.untagged t
+
     let[@inline] is_promotable ~src ~dst =
-      let[@inline] untagged = function
-        | Untagged t -> t
-        | Tagged t -> Tagged_integer.untagged t
-      in
       Integer.is_promotable ~src:(untagged src) ~dst:(untagged dst)
 
     let static_cast ~dbg ~src ~dst exp =

backend/cmm_helpers.mli

@@ -1285,7 +1285,12 @@ module Numeric : sig
 
     val nativeint : t
 
-    val static_cast : t static_cast
+    (** Gets the integral resulting from untagging the integer (if it is tagged).
+
+        E.g., you can use [static_cast ~src ~dst:(Untagged (untagged src))] to untag a
+        value of type [src]
+    *)
+    val untagged : t -> Integer.t
 
     include Integral_ops with type t := t
   end

middle_end/flambda2/to_cmm/to_cmm_primitive.ml

@@ -720,46 +720,35 @@ let binary_int_arith_primitive _env dbg (kind : K.Standard_int.t)
      | Xor -> wrap C.xor_int)
 
 let binary_int_shift_primitive _env dbg kind (op : P.int_shift_op) x y =
-  (* [kind] only applies to [x], the [y] argument is always a bare
-     register-sized integer *)
   (* See comments on [binary_int_arity_primitive], above, about sign extension
      and use of [C.low_bits]. *)
   let kind = integral_of_standard_int kind in
-  let signedness : C.Numeric.Signedness.t =
-    match op with
-    | Lsl -> C.Numeric.Integral.signedness kind
-    | Asr -> C.Numeric.Signedness.Signed
-    | Lsr -> C.Numeric.Signedness.Unsigned
-  in
-  let kind_with_signedness_of_operator =
-    C.Numeric.Integral.with_signedness kind ~signedness
+  let right_shift_kind signedness =
+    (* right shifts can operate directly on untagged small integers of the
+       correct signedness, as they do not require sign- or zero-extension after
+       the shift *)
+    C.Numeric.Integer.with_signedness
+      (C.Numeric.Integral.untagged kind)
+      ~signedness
   in
-  let operator_type =
-    C.Numeric.Integral.(with_signedness nativeint) ~signedness
-  in
-  let x =
-    C.Numeric.Integral.static_cast x ~dbg
-      ~src:kind
-      ~dst:kind_with_signedness_of_operator
+  let f, (op_kind : Cmm_helpers.Numeric.Integer.t) =
+    match op with
+    | Asr -> C.asr_int, right_shift_kind Signed
+    | Lsr -> C.lsr_int, right_shift_kind Unsigned
+    | Lsl ->
+      (* Left shifts need to be casted back from nativeint since they might
+         shift arbitrary bits into the high bits of the register. *)
+      C.lsl_int, C.Numeric.Integer.nativeint
   in
   let x =
-    C.Numeric.Integral.static_cast x ~dbg
-      ~src:kind_with_signedness_of_operator
-      ~dst:operator_type
-  in
-  let untagged_kind_with_signedness_of_operator : C.Numeric.Integral.t =
-    Untagged
-      (match kind_with_signedness_of_operator with
-        | Tagged x -> C.Numeric.Tagged_integer.untagged x
-        | Untagged x -> x)
+    C.Numeric.Integral.static_cast x ~dbg ~src:kind ~dst:(Untagged op_kind)
   in
-  let shifted, output =
-    match op with
-    | Lsl -> C.lsl_int x y dbg, operator_type
-    | Asr -> C.asr_int x y dbg, untagged_kind_with_signedness_of_operator
-    | Lsr -> C.lsr_int x y dbg, untagged_kind_with_signedness_of_operator
+  let shifted =
+    (* [kind] only applies to [x], the [y] argument is always a bare
+       register-sized integer *)
+    f x y dbg
   in
-  C.Numeric.Integral.static_cast shifted ~src:output ~dst:kind ~dbg
+  C.Numeric.Integral.static_cast shifted ~dbg ~src:(Untagged op_kind) ~dst:kind
 
 let binary_int_comp_primitive _env dbg kind cmp x y =
   match