From d7403b3aa880065c6c65f546a7a66aca87d95999 Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Fri, 19 Jul 2019 18:22:57 -0700
Subject: [PATCH] Propagate inboudns through setindex!!

---
 benchmark/map.jl   | 22 ++++++++++++++++++++++
 src/NoBang/base.jl |  4 +++-
 src/base.jl        |  4 ++--
 src/core.jl        |  2 +-
 test/Manifest.toml |  8 ++++++++
 test/Project.toml  |  2 ++
 test/__test_ir.jl  | 28 ++++++++++++++++++++++++++++
 test/test_ir.jl    |  6 ++++++
 8 files changed, 72 insertions(+), 4 deletions(-)
 create mode 100644 benchmark/map.jl
 create mode 100644 test/__test_ir.jl
 create mode 100644 test/test_ir.jl

diff --git a/benchmark/map.jl b/benchmark/map.jl
new file mode 100644
index 00000000..5b1f61aa
--- /dev/null
+++ b/benchmark/map.jl
@@ -0,0 +1,22 @@
+using BangBang
+
+
+mapbb(f, xs::Array, inbounds::Val = Val(true)) =
+    __map!!(f, similar(xs, Union{}), xs, 1, inbounds)
+
+
+@inline function __map!!(f, ys, xs, n, inbounds)
+    @simd for i in n:length(ys)
+        if inbounds === Val(true)
+            @inbounds begin
+                ys′ = setindex!!(ys, f(xs[i]), i)
+            end
+        else
+            ys′ = setindex!!(ys, f(xs[i]), i)
+        end
+        if !(ys′ isa typeof(ys))
+            return __map!!(f, ys′, xs, i + 1, inbounds)
+        end
+    end
+    return ys
+end
diff --git a/src/NoBang/base.jl b/src/NoBang/base.jl
index f2a87eb5..a168ee48 100644
--- a/src/NoBang/base.jl
+++ b/src/NoBang/base.jl
@@ -99,7 +99,9 @@ _setindex(xs::NamedTuple, value, name) = setproperty(xs, name, value)
 function _setindex(xs::AbstractArray, v, I...)
     T = promote_type(eltype(xs), typeof(v))
     ys = similar(xs, T)
-    copy!(ys, xs)
+    if eltype(xs) !== Union{}
+        copy!(ys, xs)
+    end
     ys[I...] = v
     return ys
 end
diff --git a/src/base.jl b/src/base.jl
index 6bbef4fc..560bada8 100644
--- a/src/base.jl
+++ b/src/base.jl
@@ -196,9 +196,9 @@ possible(::typeof(_splice!), xs::C, i::Any, ::S) where {C, S} =
 """
     setindex!!(collection, value, indices...) -> collection′
 """
-setindex!!(xs, v, I...) = may(_setindex!, xs, v, I...)
+Base.@propagate_inbounds setindex!!(xs, v, I...) = may(_setindex!, xs, v, I...)
 
-_setindex!(xs, v, I...) = (setindex!(xs, v, I...); xs)
+Base.@propagate_inbounds _setindex!(xs, v, I...) = (setindex!(xs, v, I...); xs)
 
 pure(::typeof(_setindex!)) = NoBang._setindex
 possible(::typeof(_setindex!), ::Union{Tuple, NamedTuple}, ::Vararg) = false
diff --git a/src/core.jl b/src/core.jl
index 9c0f4cb0..1afc9e92 100644
--- a/src/core.jl
+++ b/src/core.jl
@@ -1,4 +1,4 @@
-may(mutate, args...) =
+Base.@propagate_inbounds may(mutate, args...) =
     if possible(mutate, args...)
         mutate(args...)
     else
diff --git a/test/Manifest.toml b/test/Manifest.toml
index 1ae145b4..76e1eb1a 100644
--- a/test/Manifest.toml
+++ b/test/Manifest.toml
@@ -51,6 +51,14 @@ uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 deps = ["Random"]
 uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
 
+[[IRTest]]
+deps = ["Test"]
+git-tree-sha1 = "5f2778f64fdc9d2e29176869ba49167777ed1170"
+repo-rev = "master"
+repo-url = "https://github.com/tkf/IRTest.jl"
+uuid = "c9680f28-2a59-47f6-98e4-d41620d42cd2"
+version = "0.1.0"
+
 [[InitialValues]]
 git-tree-sha1 = "00fc3631033606a66b6eebce4b020ad21abd5243"
 repo-rev = "master"
diff --git a/test/Project.toml b/test/Project.toml
index 8cee66d0..31e5f0c1 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,7 +1,9 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+IRTest = "c9680f28-2a59-47f6-98e4-d41620d42cd2"
 InitialValues = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/__test_ir.jl b/test/__test_ir.jl
new file mode 100644
index 00000000..80f185db
--- /dev/null
+++ b/test/__test_ir.jl
@@ -0,0 +1,28 @@
+using InteractiveUtils: code_llvm
+using Test
+
+include("../benchmark/map.jl")
+
+"""
+    llvm_ir(f, args) :: String
+
+Get LLVM IR of `f(args...)` as a string.
+"""
+llvm_ir(f, args) = sprint(code_llvm, f, Base.typesof(args...))
+
+nmatches(r, s) = count(_ -> true, eachmatch(r, s))
+
+twice(x) = 2x
+
+@testset "setindex!! (@inbounds)" begin
+    @test mapbb(twice, [1:3;], Val(false)) == [2, 4, 6]
+
+    ir_noinbounds = llvm_ir(mapbb, (twice, Float64[], Val(false)))
+    if Base.JLOptions().check_bounds == 2  # --check-bounds=no
+        @test nmatches(r"fmul <4 x double>", ir_noinbounds) >= 4
+    else
+        @test nmatches(r"fmul <4 x double>", ir_noinbounds) == 0
+    end
+    ir_inbounds = llvm_ir(mapbb, (twice, Float64[], Val(true)))
+    @test nmatches(r"fmul <4 x double>", ir_inbounds) >= 4
+end
diff --git a/test/test_ir.jl b/test/test_ir.jl
new file mode 100644
index 00000000..66c08d3c
--- /dev/null
+++ b/test/test_ir.jl
@@ -0,0 +1,6 @@
+module TestIR
+
+using IRTest
+IRTest.@include("__test_ir.jl")
+
+end  # module