- не делаем аллокаций при получении valuesIdx

commit_hash:6f19d55db0d3c32c2a374ce7c4d68f9c1a4b6c24
YandexClassifieds · Nov 20, 2024 · c909497 · c909497
1 parent 7e802c7
commit c909497
Show file tree

Hide file tree

Showing 11 changed files with 2,999 additions and 118 deletions.
diff --git a/memory/generate.sc b/memory/generate.sc
@@ -23,6 +23,12 @@ def setMethod(
   else
     s"def $name(mem: Memory[L, Global], value: $tpe ${ctArgs(i, argType)}): Unit"
 
+def getMemorySegmentMethod(withRegion: Boolean)(argType: String)(name: String, i: Int, tpe: String): String =
+  if (withRegion)
+    s"def $name[R <: Global](mem: Memory[L, R] ${ctArgs(i, argType)}, size: Int)(implicit region: Region[R]): MemorySegment"
+  else
+    s"def $name(mem: Memory[L, Global] ${ctArgs(i, argType)}, size: Int): $tpe"
+
 def generateOffsetCalculation(i: Int): String = {
   var result = "vh.offset0"
   for (j <- 1 to i) {
@@ -50,9 +56,14 @@ def setMethodImpl(typ: String, i: Int): String = {
   s"mem.asJava.set(${typeLayout(typ)}, ${generateOffsetCalculation(i)}, value)"
 }
 
+def getMemorySegmentMethodImpl(typ: String, i: Int): String = {
+  s"mem.asJava.asSlice(${generateOffsetCalculation(i)}, size * ${typeLayout(typ)}.byteSize())"
+}
+
 val methods = List(
   ("get", (getMethod(withRegion = true) _, getMethodImpl _)),
   ("set", (setMethod(withRegion = true) _, setMethodImpl _)),
+  ("getMemorySegment", (getMemorySegmentMethod(withRegion = true) _, getMemorySegmentMethodImpl _))
 )
 
 def ctParams(i: Int): String =
@@ -125,7 +136,7 @@ def memoryHandle =
 import alien.memory.{Global, Layout, Memory, Region}
 import jdk.internal.vm.annotation.IntrinsicCandidate
 
-import java.lang.foreign.ValueLayout
+import java.lang.foreign.{MemorySegment, ValueLayout}
 
 object MemoryHandle {
     trait MemoryHandleSyntax {