From 774342570aeeb24fd2964e898cbdab01b1468ae2 Mon Sep 17 00:00:00 2001
From: Yoshiyuki Mineo <Yoshiyuki.Mineo@jp.sony.com>
Date: Wed, 7 May 2025 13:17:30 +0900
Subject: [PATCH] Add Compose method (#71)

---
 v2/sonyflake.go      |  11 +++++
 v2/sonyflake_test.go | 100 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 95 insertions(+), 16 deletions(-)

diff --git a/v2/sonyflake.go b/v2/sonyflake.go
index e6ab73a..244e1dd 100644
--- a/v2/sonyflake.go
+++ b/v2/sonyflake.go
@@ -252,6 +252,17 @@ func (sf *Sonyflake) ToTime(id int64) time.Time {
 	return time.Unix(0, (sf.startTime+sf.timePart(id))*sf.timeUnit)
 }
 
+// Compose creates a Sonyflake ID from its components.
+// The time parameter should be the time when the ID was generated.
+// The sequence parameter should be between 0 and 2^BitsSequence-1 (inclusive).
+// The machineID parameter should be between 0 and 2^BitsMachineID-1 (inclusive).
+func (sf *Sonyflake) Compose(t time.Time, sequence, machineID int) int64 {
+	elapsedTime := sf.toInternalTime(t.UTC()) - sf.startTime
+	return elapsedTime<<(sf.bitsSequence+sf.bitsMachine) |
+		int64(sequence)<<sf.bitsMachine |
+		int64(machineID)
+}
+
 // Decompose returns a set of Sonyflake ID parts.
 func (sf *Sonyflake) Decompose(id int64) map[string]int64 {
 	time := sf.timePart(id)
diff --git a/v2/sonyflake_test.go b/v2/sonyflake_test.go
index 85e747c..30d97b0 100644
--- a/v2/sonyflake_test.go
+++ b/v2/sonyflake_test.go
@@ -253,22 +253,6 @@ func TestNextID_ReturnsError(t *testing.T) {
 	}
 }
 
-func TestToTime(t *testing.T) {
-	start := time.Now()
-	sf := newSonyflake(t, Settings{
-		TimeUnit:  time.Millisecond,
-		StartTime: start,
-	})
-
-	id := nextID(t, sf)
-
-	tm := sf.ToTime(id)
-	diff := tm.Sub(start)
-	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
-		t.Errorf("unexpected time: %v", tm)
-	}
-}
-
 func TestPrivateIPv4(t *testing.T) {
 	testCases := []struct {
 		description    string
@@ -351,3 +335,87 @@ func TestLower16BitPrivateIP(t *testing.T) {
 		})
 	}
 }
+
+func TestToTime(t *testing.T) {
+	start := time.Now()
+	sf := newSonyflake(t, Settings{
+		TimeUnit:  time.Millisecond,
+		StartTime: start,
+	})
+
+	id := nextID(t, sf)
+
+	tm := sf.ToTime(id)
+	diff := tm.Sub(start)
+	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
+		t.Errorf("unexpected time: %v", tm)
+	}
+}
+
+func TestComposeAndDecompose(t *testing.T) {
+	now := time.Now()
+	sf := newSonyflake(t, Settings{
+		TimeUnit:  time.Millisecond,
+		StartTime: now,
+	})
+
+	testCases := []struct {
+		name      string
+		time      time.Time
+		sequence  int
+		machineID int
+	}{
+		{
+			name:      "zero values",
+			time:      now,
+			sequence:  0,
+			machineID: 0,
+		},
+		{
+			name:      "max sequence",
+			time:      now,
+			sequence:  1<<sf.bitsSequence - 1,
+			machineID: 0,
+		},
+		{
+			name:      "max machine id",
+			time:      now,
+			sequence:  0,
+			machineID: 1<<sf.bitsMachine - 1,
+		},
+		{
+			name:      "future time",
+			time:      now.Add(time.Hour),
+			sequence:  0,
+			machineID: 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			id := sf.Compose(tc.time, tc.sequence, tc.machineID)
+			parts := sf.Decompose(id)
+
+			// Verify time part
+			expectedTime := sf.toInternalTime(tc.time.UTC()) - sf.startTime
+			if parts["time"] != expectedTime {
+				t.Errorf("time mismatch: got %d, want %d", parts["time"], expectedTime)
+			}
+
+			// Verify sequence part
+			if parts["sequence"] != int64(tc.sequence) {
+				t.Errorf("sequence mismatch: got %d, want %d", parts["sequence"], tc.sequence)
+			}
+
+			// Verify machine id part
+			if parts["machine"] != int64(tc.machineID) {
+				t.Errorf("machine id mismatch: got %d, want %d", parts["machine"], tc.machineID)
+			}
+
+			// Verify id part
+			if parts["id"] != id {
+				t.Errorf("id mismatch: got %d, want %d", parts["id"], id)
+			}
+		})
+	}
+}