Use MacroTools in @maybe_grow

ITensor · May 13, 2024 · 625fa27 · 625fa27
1 parent 41ef815
commit 625fa27
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml
@@ -18,6 +18,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -34,20 +35,20 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
 NDTensorsAMDGPUExt = "AMDGPU"
 NDTensorsCUDAExt = "CUDA"
-NDTensorscuTENSORExt = "cuTENSOR"
 NDTensorsHDF5Ext = "HDF5"
 NDTensorsMetalExt = "Metal"
 NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"
+NDTensorscuTENSORExt = "cuTENSOR"
 
 [compat]
 Accessors = "0.1.33"
@@ -65,6 +66,7 @@ HDF5 = "0.14, 0.15, 0.16, 0.17"
 HalfIntegers = "1"
 InlineStrings = "1"
 LinearAlgebra = "1.6"
+MacroTools = "0.5"
 MappedArrays = "0.4"
 PackageExtensionCompat = "1"
 Random = "1.6"

diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl
@@ -1,5 +1,6 @@
 using Accessors: @set
 using Dictionaries: Dictionary, set!
+using MacroTools: @capture
 using ..SparseArrayInterface:
   SparseArrayInterface, AbstractSparseArray, getindex_zero_function
 
@@ -124,8 +125,6 @@ macro maybe_grow(ex)
       "@maybe_grow must be used with setindex! syntax (as @maybe_grow a[i,j,...] = value)"
     )
   end
-  arr_name = esc(ex.args[1].args[1])
-  indices = esc(ex.args[1].args[2:end])
-  value = esc(ex.args[2])
-  return :(setindex_maybe_grow!($arr_name, $value, $indices...))
+  @capture(ex, array_[indices__] = value_)
+  return :(setindex_maybe_grow!($(esc(array)), $value, $indices...))
 end