Skip to content

Commit 3622c49

Browse files
andreypfauandreypfau
committed
Pruned branch cells
Signed-off-by: andreypfau <andreypfau@ton.org>
1 parent 6abefd9 commit 3622c49

File tree

10 files changed

+141
-13
lines changed

10 files changed

+141
-13
lines changed

‎bitstring/src/BitString.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,6 @@ public class BitString internal constructor(
125125
internal fun getBackingArrayReference(): ByteArray = data
126126

127127
public companion object {
128-
internal val EMPTY = BitString(ByteArray(0), 0, 0)
128+
public val EMPTY: BitString = BitString(ByteArray(0), 0, 0)
129129
}
130130
}

‎bitstring/src/internal/BitsOperations.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -482,6 +482,9 @@ internal fun storeIntIntoByteArray(
482482
v = v shl (32 - topBits)
483483
}
484484

485+
val destOffset = destOffset + (bitOffset ushr 3)
486+
val bitOffset = bitOffset and 7
487+
485488
// Fast path: byte-aligned start and length is a multiple of whole bytes.
486489
if (bitOffset == 0 && (topBits and 7) == 0) {
487490
val byteLen = topBits ushr 3

‎cell/build.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ kotlin {
1818
api(projects.tonSdkBigint)
1919
api(libs.kotlinx.io.bytestring)
2020
implementation(libs.serialization.core)
21+
implementation(libs.atomicfu)
2122
}
2223
}
2324
}

‎cell/src/Cell.kt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.ton.sdk.cell
22

3+
import org.ton.sdk.bitstring.BitString
34
import org.ton.sdk.cell.internal.EmptyCell
45
import org.ton.sdk.crypto.HashBytes
56
import kotlin.jvm.JvmStatic
@@ -24,7 +25,7 @@ public interface Cell {
2425
* Returns this cell as a virtualized cell, so that all hashes
2526
* and depths will have an offset.
2627
*/
27-
public fun virtualize(offset: Int): Cell
28+
public fun virtualize(offset: Int = 1): Cell
2829

2930
/**
3031
* Returns the hash of the current cell for the specified level.
@@ -58,5 +59,7 @@ public interface Cell {
5859
}
5960

6061
public interface LoadedCell : Cell {
62+
public val bits: BitString
63+
6164
public fun reference(index: Int): Cell?
6265
}

‎cell/src/CellBuilder.kt‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,12 @@ import org.ton.sdk.bitstring.unsafe.UnsafeBitStringOperations
99
import org.ton.sdk.cell.internal.DataCell
1010
import org.ton.sdk.cell.internal.EmptyCell
1111
import org.ton.sdk.cell.internal.LibraryCell
12+
import org.ton.sdk.cell.internal.PrunedCell
1213
import org.ton.sdk.crypto.HashBytes
1314
import org.ton.sdk.crypto.Sha256
1415
import kotlin.experimental.and
1516
import kotlin.experimental.or
17+
import kotlin.jvm.JvmStatic
1618
import kotlin.math.max
1719

1820
public class CellBuilder() {
@@ -157,7 +159,7 @@ public class CellBuilder() {
157159
}
158160

159161
CellType.LIBRARY_REFERENCE -> LibraryCell(descriptor, hashes[0])
160-
CellType.PRUNED_BRANCH -> TODO()
162+
CellType.PRUNED_BRANCH -> PrunedCell(descriptor, hashes[0], bitString)
161163
CellType.MERKLE_PROOF,
162164
CellType.MERKLE_UPDATE -> DataCell(
163165
descriptor,
@@ -232,6 +234,44 @@ public class CellBuilder() {
232234
}
233235
return descriptor
234236
}
237+
238+
public companion object {
239+
@JvmStatic
240+
public fun createPrunedBranch(
241+
cell: Cell,
242+
newLevel: Int,
243+
): Cell {
244+
return createPrunedBranch(cell, newLevel, CellContext.EMPTY)
245+
}
246+
247+
@JvmStatic
248+
public fun createPrunedBranch(
249+
cell: Cell,
250+
newLevel: Int,
251+
cellContext: CellContext = CellContext.EMPTY
252+
): Cell {
253+
if (cell is LoadedCell && cell.descriptor.referenceCount == 0) {
254+
return cell
255+
}
256+
var cellLevelMask = cell.levelMask
257+
val levelMask = LevelMask(cellLevelMask.mask or newLevel)
258+
val builder = CellBuilder()
259+
builder.storeUInt(CellType.PRUNED_BRANCH.value, 8)
260+
builder.storeUInt(levelMask.mask, 8)
261+
262+
// Only write levels lower than the new level.
263+
cellLevelMask = LevelMask(cellLevelMask.mask and (newLevel - 1))
264+
265+
cellLevelMask.forEach { level ->
266+
builder.store(cell.hash(level))
267+
}
268+
cellLevelMask.forEach { level ->
269+
builder.storeUInt(cell.depth(level), 16)
270+
}
271+
272+
return builder.build(isExotic = true)
273+
}
274+
}
235275
}
236276

237277
public operator fun CellBuilder.plusAssign(builder: CellBuilder) {

‎cell/src/LevelMask.kt‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import kotlin.jvm.JvmStatic
77
*/
88
public class LevelMask(
99
mask: Int = 0
10-
) {
10+
) : Iterable<Int> {
1111
public val mask: Int = mask and 0b111
1212

1313
/**
@@ -54,6 +54,16 @@ public class LevelMask(
5454
public fun virtualize(offset: Int = 1): LevelMask =
5555
LevelMask(mask ushr offset)
5656

57+
override fun iterator(): Iterator<Int> = iterator {
58+
var currentMask = 1 or (mask shl 1)
59+
while (currentMask != 0) {
60+
val levelMask = currentMask and (currentMask.inv() + 1)
61+
val level = levelMask.countTrailingZeroBits()
62+
yield(level)
63+
currentMask = currentMask and levelMask.inv()
64+
}
65+
}
66+
5767
override fun equals(other: Any?): Boolean {
5868
if (this === other) return true
5969
if (other !is LevelMask) return false

‎cell/src/internal/DataCell.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import org.ton.sdk.crypto.HashBytes
88

99
internal class DataCell(
1010
override val descriptor: CellDescriptor,
11-
val bits: BitString,
11+
override val bits: BitString,
1212
private val references: List<Cell>,
1313
private val hashes: Array<HashBytes>,
1414
private val depths: IntArray

‎cell/src/internal/EmptyCell.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.ton.sdk.cell.internal
22

3+
import org.ton.sdk.bitstring.BitString
34
import org.ton.sdk.cell.Cell
45
import org.ton.sdk.cell.CellDescriptor
56
import org.ton.sdk.cell.LoadedCell
@@ -14,6 +15,8 @@ internal object EmptyCell : LoadedCell {
1415

1516
override fun depth(level: Int): Int = 0
1617

18+
override val bits: BitString = BitString.EMPTY
19+
1720
override fun reference(index: Int): Cell? = null
1821

1922
val EMPTY_CELL_HASH = HashBytes(

‎cell/src/internal/PrunedCell.kt‎

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,71 @@
11
package org.ton.sdk.cell.internal
22

3+
import kotlinx.io.bytestring.unsafe.UnsafeByteStringApi
4+
import kotlinx.io.bytestring.unsafe.UnsafeByteStringOperations
5+
import org.ton.sdk.bitstring.BitString
6+
import org.ton.sdk.bitstring.unsafe.UnsafeBitStringApi
7+
import org.ton.sdk.bitstring.unsafe.UnsafeBitStringOperations
38
import org.ton.sdk.cell.Cell
49
import org.ton.sdk.cell.CellDescriptor
510
import org.ton.sdk.crypto.HashBytes
611

712
internal class PrunedCell(
813
override val descriptor: CellDescriptor,
914
val hash: HashBytes,
10-
val data: ByteArray
15+
val data: BitString
1116
) : Cell {
12-
override fun virtualize(offset: Int): Cell {
13-
TODO("Not yet implemented")
14-
}
17+
override fun virtualize(offset: Int): Cell = VirtualCell(this, offset)
1518

19+
@OptIn(UnsafeByteStringApi::class, UnsafeBitStringApi::class)
1620
override fun hash(level: Int): HashBytes {
1721
val hashIndex = descriptor.levelMask.apply(level).hashIndex
18-
if (hashIndex == this.level) {
22+
val cellLevel = this.level
23+
if (hashIndex == cellLevel) {
1924
return hash
2025
}
21-
22-
TODO()
26+
val offset = 2 + hashIndex * 32
27+
return UnsafeBitStringOperations.withByteArrayUnsafe(data) {
28+
HashBytes(
29+
UnsafeByteStringOperations.wrapUnsafe(it.copyOfRange(offset, offset + 32))
30+
)
31+
}
2332
}
2433

34+
@Suppress("OPT_IN_USAGE")
2535
override fun depth(level: Int): Int {
26-
TODO("Not yet implemented")
36+
val hashIndex = descriptor.levelMask.apply(level).hashIndex
37+
val cellLevel = this.level
38+
if (hashIndex == cellLevel) {
39+
return 0
40+
}
41+
val offset = 2 + cellLevel * 32 + hashIndex * 2
42+
return UnsafeBitStringOperations.withByteArrayUnsafe(data) {
43+
(it[offset].toInt() and 0xFF shl 8) or
44+
it[offset + 1].toInt() and 0xFF
45+
}
46+
}
47+
48+
override fun toString(): String {
49+
return "PrunedCell(descriptor=$descriptor, hash=$hash, data=$data)"
50+
}
51+
52+
override fun equals(other: Any?): Boolean {
53+
if (this === other) return true
54+
if (other == null || this::class != other::class) return false
55+
56+
other as PrunedCell
57+
58+
if (descriptor != other.descriptor) return false
59+
if (hash != other.hash) return false
60+
if (data != other.data) return false
61+
62+
return true
63+
}
64+
65+
override fun hashCode(): Int {
66+
var result = descriptor.hashCode()
67+
result = 31 * result + hash.hashCode()
68+
result = 31 * result + data.hashCode()
69+
return result
2770
}
2871
}

‎cell/test/PrunedBranchTest.kt‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
import org.ton.sdk.cell.Cell
2+
import org.ton.sdk.cell.CellBuilder
3+
import kotlin.test.Test
4+
import kotlin.test.assertEquals
5+
6+
class PrunedBranchTest {
7+
@Test
8+
fun correctPrunedBranch() {
9+
val cell = CellBuilder()
10+
.storeULong(0xdeafbeaf123123, 128)
11+
.storeReference(Cell.EMPTY)
12+
.build()
13+
14+
val prunedCell = CellBuilder.createPrunedBranch(cell, 1)
15+
assertEquals(cell.hash(), prunedCell.hash(0))
16+
assertEquals(cell.depth(), prunedCell.depth(0))
17+
18+
val virtualCell = cell.virtualize()
19+
assertEquals(cell.hash(), virtualCell.hash())
20+
assertEquals(cell.depth(3), virtualCell.depth(3))
21+
22+
val virtualPrunedBranch = CellBuilder.createPrunedBranch(virtualCell, 1)
23+
assertEquals(prunedCell, virtualPrunedBranch)
24+
}
25+
}

0 commit comments

Comments
 (0)