From dc9a0b31c0ed67f6e52cc50002b1a3f0e43d5248 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 24 May 2024 19:55:21 +0800
Subject: [PATCH] lntypes: add new units `WeightUnit` and `VByte`

---
 lntypes/txsize.go      | 42 ++++++++++++++++++++++++++++++++++++++++++
 lntypes/txsize_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100644 lntypes/txsize.go
 create mode 100644 lntypes/txsize_test.go

diff --git a/lntypes/txsize.go b/lntypes/txsize.go
new file mode 100644
index 000000000..0618bfc51
--- /dev/null
+++ b/lntypes/txsize.go
@@ -0,0 +1,42 @@
+package lntypes
+
+import (
+	"fmt"
+	"math"
+)
+
+// WeightUnit defines a unit to express the transaction size. One weight unit
+// is 1/4_000_000 of the max block size. The tx weight is calculated using
+// `Base tx size * 3 + Total tx size`.
+//   - Base tx size is size of the transaction serialized without the witness
+//     data.
+//   - Total tx size is the transaction size in bytes serialized according
+//     #BIP144.
+type WeightUnit uint64
+
+// ToVB converts a value expressed in weight units to virtual bytes.
+func (wu WeightUnit) ToVB() VByte {
+	// According to BIP141: Virtual transaction size is defined as
+	// Transaction weight / 4 (rounded up to the next integer).
+	return VByte(math.Ceil(float64(wu) / 4))
+}
+
+// String returns the string representation of the weight unit.
+func (wu WeightUnit) String() string {
+	return fmt.Sprintf("%d wu", wu)
+}
+
+// VByte defines a unit to express the transaction size. One virtual byte is
+// 1/4th of a weight unit. The tx virtual bytes is calculated using `TxWeight /
+// 4`.
+type VByte uint64
+
+// ToWU converts a value expressed in virtual bytes to weight units.
+func (vb VByte) ToWU() WeightUnit {
+	return WeightUnit(vb * 4)
+}
+
+// String returns the string representation of the virtual byte.
+func (vb VByte) String() string {
+	return fmt.Sprintf("%d vb", vb)
+}
diff --git a/lntypes/txsize_test.go b/lntypes/txsize_test.go
new file mode 100644
index 000000000..a68a45b23
--- /dev/null
+++ b/lntypes/txsize_test.go
@@ -0,0 +1,28 @@
+package lntypes
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestTxSizeUnit tests the conversion of tx size to different units.
+func TestTxSizeUnit(t *testing.T) {
+	t.Parallel()
+
+	// Test normal conversion from 100wu to 25vb.
+	wu := WeightUnit(100)
+	vb := VByte(25)
+	require.Equal(t, vb, wu.ToVB(), "wu -> vb conversion "+
+		"failed: want %v, got %v", vb, wu.ToVB())
+	require.Equal(t, wu, vb.ToWU(), "vb -> wu conversion "+
+		"failed: want %v, got %v", wu, vb.ToWU())
+
+	// Test rounding up conversion from 99wu to 25vb.
+	wu = WeightUnit(99)
+	vb = VByte(25)
+	require.Equal(t, vb, wu.ToVB(), "wu -> vb conversion "+
+		"failed: want %v, got %v", vb, wu.ToVB())
+	require.Equal(t, WeightUnit(100), vb.ToWU(), "vb -> wu conversion "+
+		"failed: want %v, got %v", 100, vb.ToWU())
+}