Fix inc. ranges with start == end and negative step

src/Futhark/Internalise/Exps.hs

@@ -184,22 +184,19 @@ internaliseAppExp desc _ (E.Index e idxs loc) = do
         v_t <- lookupType v
         pure $ I.BasicOp $ I.Index v $ fullSlice v_t idxs'
   certifying cs $ mapM (letSubExp desc <=< index) vs
-internaliseAppExp desc _ (E.Range start maybe_second end loc) = do
+internaliseAppExp desc _ (E.Range start maybe_second inclusiveness loc) = do
   start' <- internaliseExp1 "range_start" start
-  end' <- internaliseExp1 "range_end" $ case end of
+  end' <- internaliseExp1 "range_end" $ case inclusiveness of
     DownToExclusive e -> e
     ToInclusive e -> e
     UpToExclusive e -> e
   maybe_second' <-
     traverse (internaliseExp1 "range_second") maybe_second
 
   -- Construct an error message in case the range is invalid.
-  let conv = case E.typeOf start of
-        E.Scalar (E.Prim (E.Unsigned _)) -> asIntZ Int64
-        _ -> asIntS Int64
-  start'_i64 <- conv start'
-  end'_i64 <- conv end'
-  maybe_second'_i64 <- traverse conv maybe_second'
+  start'_i64 <- asIntS Int64 start'
+  end'_i64 <- asIntS Int64 end'
+  maybe_second'_i64 <- traverse (asIntS Int64) maybe_second'
   let errmsg =
         errorMsg $
           ["Range "]
@@ -208,104 +205,83 @@ internaliseAppExp desc _ (E.Range start maybe_second end loc) = do
                    Nothing -> []
                    Just second_i64 -> ["..", ErrorVal int64 second_i64]
                )
-            ++ ( case end of
+            ++ ( case inclusiveness of
                    DownToExclusive {} -> ["..>"]
                    ToInclusive {} -> ["..."]
                    UpToExclusive {} -> ["..<"]
                )
             ++ [ErrorVal int64 end'_i64, " is invalid."]
 
-  (it, le_op, lt_op) <-
-    case E.typeOf start of
-      E.Scalar (E.Prim (E.Signed it)) -> pure (it, CmpSle it, CmpSlt it)
-      E.Scalar (E.Prim (E.Unsigned it)) -> pure (it, CmpUle it, CmpUlt it)
-      start_t -> error $ "Start value in range has type " ++ prettyString start_t
-
-  let one = intConst it 1
-      negone = intConst it (-1)
-      default_step = case end of
-        DownToExclusive {} -> negone
-        ToInclusive {} -> one
-        UpToExclusive {} -> one
-
-  (step, step_zero) <- case maybe_second' of
-    Just second' -> do
-      subtracted_step <-
-        letSubExp "subtracted_step" $
+  let it = case E.typeOf start of
+        E.Scalar (E.Prim (E.Signed start_t)) -> start_t
+        start_t -> error $ "Start value in range has type " ++ prettyString start_t
+  let lt_op = CmpSlt it
+
+  let zero = intConst it 0
+
+  step <- case maybe_second' of
+    Just second' ->
+      letSubExp "subtracted_step" $
+        I.BasicOp $
+          I.BinOp (I.Sub it I.OverflowWrap) second' start'
+    Nothing -> pure $
+      intConst it $ -- use default step of 1 or -1.
+        case inclusiveness of
+          DownToExclusive {} -> -1
+          _ -> 1
+
+  step_zero <-
+    letSubExp "step_zero" $
+      I.BasicOp $
+        I.CmpOp (I.CmpEq $ IntType it) step zero
+
+  distance <- do
+    difference <-
+      letSubExp "difference" $
+        I.BasicOp $
+          I.BinOp (Sub it I.OverflowWrap) end' start'
+    distance_exclusive_i64 <-
+      asIntS Int64
+        =<< letSubExp
+          "distance_exclusive"
+          (I.BasicOp $ I.UnOp (Abs it) difference)
+
+    case inclusiveness of
+      ToInclusive {} ->
+        letSubExp "distance_inclusive" $
           I.BasicOp $
-            I.BinOp (I.Sub it I.OverflowWrap) second' start'
-      step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second'
-      pure (subtracted_step, step_zero)
-    Nothing ->
-      pure (default_step, constant False)
+            I.BinOp
+              (Add Int64 I.OverflowWrap)
+              distance_exclusive_i64
+              (intConst Int64 1)
+      _ -> pure distance_exclusive_i64
 
-  step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step
-  step_sign_i64 <- asIntS Int64 step_sign
+  downwards <-
+    letSubExp "downwards" $
+      I.BasicOp $
+        I.CmpOp lt_op step zero
+  step_wrong_dir <- case inclusiveness of
+    DownToExclusive {} -> letSubExp "upwards" $ I.BasicOp $ I.UnOp (I.Neg I.Bool) downwards
+    UpToExclusive {} -> pure downwards
+    ToInclusive {} -> pure $ constant False
 
   bounds_invalid_downwards <-
     letSubExp "bounds_invalid_downwards" $
       I.BasicOp $
-        I.CmpOp le_op start' end'
+        I.CmpOp lt_op start' end'
   bounds_invalid_upwards <-
     letSubExp "bounds_invalid_upwards" $
       I.BasicOp $
         I.CmpOp lt_op end' start'
-
-  (distance, step_wrong_dir, bounds_invalid) <- case end of
-    DownToExclusive {} -> do
-      step_wrong_dir <-
-        letSubExp "step_wrong_dir" $
-          I.BasicOp $
-            I.CmpOp (I.CmpEq $ IntType it) step_sign one
-      distance <-
-        letSubExp "distance" $
-          I.BasicOp $
-            I.BinOp (Sub it I.OverflowWrap) start' end'
-      distance_i64 <- asIntS Int64 distance
-      pure (distance_i64, step_wrong_dir, bounds_invalid_downwards)
-    UpToExclusive {} -> do
-      step_wrong_dir <-
-        letSubExp "step_wrong_dir" $
-          I.BasicOp $
-            I.CmpOp (I.CmpEq $ IntType it) step_sign negone
-      distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start'
-      distance_i64 <- asIntS Int64 distance
-      pure (distance_i64, step_wrong_dir, bounds_invalid_upwards)
-    ToInclusive {} -> do
-      downwards <-
-        letSubExp "downwards" $
-          I.BasicOp $
-            I.CmpOp (I.CmpEq $ IntType it) step_sign negone
-      distance_downwards_exclusive <-
-        letSubExp "distance_downwards_exclusive" $
-          I.BasicOp $
-            I.BinOp (Sub it I.OverflowWrap) start' end'
-      distance_upwards_exclusive <-
-        letSubExp "distance_upwards_exclusive" $
-          I.BasicOp $
-            I.BinOp (Sub it I.OverflowWrap) end' start'
-
-      bounds_invalid <-
-        letSubExp "bounds_invalid"
-          =<< eIf
-            (eSubExp downwards)
-            (resultBodyM [bounds_invalid_downwards])
-            (resultBodyM [bounds_invalid_upwards])
-      distance_exclusive <-
-        letSubExp "distance_exclusive"
-          =<< eIf
-            (eSubExp downwards)
-            (resultBodyM [distance_downwards_exclusive])
-            (resultBodyM [distance_upwards_exclusive])
-      distance_exclusive_i64 <- asIntS Int64 distance_exclusive
-      distance <-
-        letSubExp "distance" $
-          I.BasicOp $
-            I.BinOp
-              (Add Int64 I.OverflowWrap)
-              distance_exclusive_i64
-              (intConst Int64 1)
-      pure (distance, constant False, bounds_invalid)
+  bounds_invalid <- case inclusiveness of
+    DownToExclusive {} -> pure bounds_invalid_downwards
+    UpToExclusive {} -> pure bounds_invalid_upwards
+    ToInclusive {} ->
+      letSubExp "bounds_invalid"
+        =<< eIf
+          (eSubExp downwards)
+          (resultBodyM [bounds_invalid_downwards])
+          (resultBodyM [bounds_invalid_upwards])
 
   step_invalid <-
     letSubExp "step_invalid" $
@@ -323,7 +299,7 @@ internaliseAppExp desc _ (E.Range start maybe_second end loc) = do
   pos_step <-
     letSubExp "pos_step" $
       I.BasicOp $
-        I.BinOp (Mul Int64 I.OverflowWrap) step_i64 step_sign_i64
+        I.UnOp (Abs Int64) step_i64
 
   num_elems <-
     certifying cs $

tests/range0.fut

@@ -5,24 +5,20 @@
 -- input { 1 5 } output { [1i32, 2i32, 3i32, 4i32, 5i32] }
 -- input { 5 1 } error: 5...1
 -- input { 5 0 } error: 5...0
--- input { 0 5 } output { [0i32, 1i32, 2i32,
--- 3i32, 4i32, 5i32] }
 
 -- ==
 -- entry: test1
 -- input { 0 5 } output { [0i32, 1i32, 2i32, 3i32, 4i32] }
 -- input { 1 5 } output { [1i32, 2i32, 3i32, 4i32] }
 -- input { 5 1 } error: 5..<1
 -- input { 5 0 } error: 5..<0
--- input { 0 5 } output { [0i32, 1i32, 2i32, 3i32, 4i32] }
 
 -- ==
 -- entry: test2
 -- input { 0 5 } error: 0..>5
 -- input { 1 5 } error: 1..>5
 -- input { 5 1 } output { [5i32, 4i32, 3i32, 2i32] }
 -- input { 5 0 } output { [5i32, 4i32, 3i32, 2i32, 1i32] }
--- input { 0 5 } error: 0..>5
 
 -- ==
 -- entry: test3
@@ -31,6 +27,9 @@
 -- input { 5 4 1 } output { [5i32, 4i32, 3i32, 2i32, 1i32] }
 -- input { 5 0 0 } output { [5i32, 0i32] }
 -- input { 0 2 5 } output { [0i32, 2i32, 4i32] }
+-- input { 1 1 5 } error: 1..1...5
+-- input { 1 0 1 } output { [1i32] }
+-- input { 1 2 1 } output { [1i32] }
 
 -- ==
 -- entry: test4