From e58c6171d11d9feb1ec1aae9fd938fd4d77e1cf2 Mon Sep 17 00:00:00 2001
From: ACEsuit <christohortner@gmail.com>
Date: Fri, 31 May 2024 16:47:33 -0700
Subject: [PATCH] cleanup

---
 src/WithAlloc.jl | 39 ++++++---------------------------------
 test/test1.jl    | 46 ----------------------------------------------
 2 files changed, 6 insertions(+), 79 deletions(-)

diff --git a/src/WithAlloc.jl b/src/WithAlloc.jl
index c97f4c4..d3d898b 100644
--- a/src/WithAlloc.jl
+++ b/src/WithAlloc.jl
@@ -5,28 +5,6 @@ export whatalloc, @withalloc
 
 function whatalloc end 
 
-@inline function _bumper_alloc(allocinfo::Tuple{<: Type, Vararg{Int, N}}) where {N}
-   (Bumper.alloc!(Bumper.default_buffer(), allocinfo...), )
-end
-
-@inline function _bumper_alloc(allocinfo)
-   map( a -> _bumper_alloc(a)[1], allocinfo )
-end
-
-# macro withalloc(ex)
-#    fncall = esc(ex.args[1])
-#    args = esc.(ex.args[2:end])
-#    quote
-#       # not sure why this isn't working ... 
-#       # whatalloc($fncall, $(args...))
-#       let 
-#          allocinfo = whatalloc($fncall, $(args...), )
-#          storobj = _bumper_alloc(allocinfo)
-#          $(fncall)(storobj..., $(args...), )
-#       end
-#    end
-# end
-
 macro withalloc(ex)
    esc_args = esc.(ex.args)
    quote
@@ -34,29 +12,24 @@ macro withalloc(ex)
    end
 end
 
-
-# For some reason that I don't understand the following implementation is allocating 
-# The @generated implementation below is to get around this. 
-# @inline function withalloc(fncall, args...) 
-#    allocinfo = whatalloc(fncall, args..., )
-#    storobj = _bumper_alloc(allocinfo)
-#    fncall(storobj..., args..., )
-# end 
-
 @inline function withalloc(fncall, args...)
    allocinfo = whatalloc(fncall, args...)
    _genwithalloc(allocinfo, fncall, args...) 
 end
 
+@inline function _bumper_alloc(allocinfo::Tuple{<: Type, Vararg{Int, N}}) where {N}
+   Bumper.alloc!(Bumper.default_buffer(), allocinfo...)
+end
+
 @inline @generated function _genwithalloc(allocinfo::TT, fncall, args...)  where {TT <: Tuple}
    code = Expr[] 
    LEN = length(TT.types) 
    if TT.types[1] <: Tuple 
       for i in 1:LEN
-         push!(code, Meta.parse("tmp$i = _bumper_alloc(allocinfo[$i])[1]"))
+         push!(code, Meta.parse("tmp$i = _bumper_alloc(allocinfo[$i])"))
       end
    else 
-      push!(code, Meta.parse("tmp1 = _bumper_alloc(allocinfo)[1]"))
+      push!(code, Meta.parse("tmp1 = _bumper_alloc(allocinfo)"))
       LEN = 1 
    end
    push!(code, Meta.parse("fncall($(join(["tmp$i, " for i in 1:LEN])) args...)"))
diff --git a/test/test1.jl b/test/test1.jl
index 1b7dd44..940d316 100644
--- a/test/test1.jl
+++ b/test/test1.jl
@@ -165,49 +165,3 @@ nalloc3 = let
 end
 
 @test nalloc3 == 0
-
-
-
-
-## Reproduce #1 
-
-using WithAlloc, LinearAlgebra, Bumper, BenchmarkTools 
-
-mymul2!(A1, A2, B, C, D) = mul!(A1, B, C), mul!(A2, B, D)
-
-function WithAlloc.whatalloc(::typeof(mymul2!), B, C, D) 
-   T1 = promote_type(eltype(B), eltype(C))
-   T2 = promote_type(eltype(B), eltype(D))
-   return ( (T1, size(B, 1), size(C, 2)), 
-            (T2, size(B, 1), size(D, 2)) )
-end
-
-function alloctest1(B, C, D) 
-   @no_escape begin 
-      a1, a2 = WithAlloc.whatalloc(mymul2!, B, C, D)
-      A1 = @alloc(a1...)
-      A2 = @alloc(a2...)
-      mymul2!(A1, A2, B, C, D) 
-      sum(A1) + sum(A2)
-   end
-end 
-
-function alloctest2(B, C, D) 
-   @no_escape begin 
-      A1, A2 = @withalloc mymul2!(B, C, D) 
-      sum(A1) + sum(A2)
-   end
-end 
-
-function alloctest3(B, C, D) 
-   @no_escape begin 
-      A1, A2 = WithAlloc.withalloc(mymul2!, B, C, D) 
-      sum(A1) + sum(A2)
-   end
-end 
-
-
-B = randn(5,10); C = randn(10, 3); D = randn(10, 5)
-@btime alloctest1($B, $C, $D)
-@btime alloctest2($B, $C, $D)
-@btime alloctest3($B, $C, $D)