Don't introduce scope in at-allowscalar.

Author: maleadt

Project.toml

@@ -15,7 +15,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 Adapt = "2.0, 3.0"
-GPUArraysCore = "= 0.1.4"
+GPUArraysCore = "= 0.1.5"
 LLVM = "3.9, 4, 5"
 Reexport = "1"
 julia = "1.6"

lib/GPUArraysCore/Project.toml

@@ -1,7 +1,7 @@
 name = "GPUArraysCore"
 uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
 authors = ["Tim Besard <[email protected]>"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/GPUArraysCore/src/GPUArraysCore.jl

@@ -110,6 +110,16 @@ function assertscalar(op = "operation")
     return
 end
 
+# Like a try-finally block, except without introducing the try scope
+# NOTE: This is deprecated and should not be used from user logic. A proper solution to
+# this problem will be introduced in https://github.com/JuliaLang/julia/pull/39217
+macro __tryfinally(ex, fin)
+    Expr(:tryfinally,
+       :($(esc(ex))),
+       :($(esc(fin)))
+       )
+end
+
 """
     @allowscalar() begin
         # code that can use scalar indexing
@@ -121,9 +131,11 @@ See also: [`allowscalar`](@ref).
 """
 macro allowscalar(ex)
     quote
-        task_local_storage(:ScalarIndexing, ScalarAllowed) do
-            $(esc(ex))
-        end
+        local tls_value = get(task_local_storage(), :ScalarIndexing, nothing)
+        task_local_storage(:ScalarIndexing, ScalarAllowed)
+        @__tryfinally($(esc(ex)),
+                      isnothing(tls_value) ? delete!(task_local_storage(), :ScalarIndexing)
+                                           : task_local_storage(:ScalarIndexing, tls_value))
     end
 end
 

test/testsuite/indexing.jl

@@ -17,6 +17,9 @@
         end
 
         @test_throws ErrorException x[]
+
+        @allowscalar y = 42
+        @test y == 42
     end
 
     @allowscalar @testset "getindex with $T" for T in eltypes