From ca280954f4f0e85b4b3ee7b6780c8b8273d23634 Mon Sep 17 00:00:00 2001
From: Timur Demin <me@tdem.in>
Date: Sat, 8 Aug 2020 20:57:23 +0500
Subject: [PATCH] Vendor IP address calculation function

The version of AddrForNodeID we vendor preallocates memory during
IP address calculation, and has a slight performance win.
---
 crypto.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 main.go   |  3 +--
 2 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 crypto.go

diff --git a/crypto.go b/crypto.go
new file mode 100644
index 0000000..35a8af1
--- /dev/null
+++ b/crypto.go
@@ -0,0 +1,47 @@
+package main
+
+import (
+	"github.com/yggdrasil-network/yggdrasil-go/src/address"
+	"github.com/yggdrasil-network/yggdrasil-go/src/crypto"
+)
+
+// AddrForNodeID is a variant of Yggdrasil's src/address.AddrForNodeID that
+// might be slightly optimized for performance.
+//
+// This function is a modded variant of address.AddrForNodeID from Yggdrasil.
+// See src/address/address.go@78b5f88e4bb734d0dd6a138ff08d34ca39dcaea3
+func AddrForNodeID(nid *crypto.NodeID) *address.Address {
+	// 128 bit address, begins with GetPrefix(), with last bit set to 0
+	// (indicates an address). Next 7 bits, interpreted as a uint, are the count
+	// of leading 1s in the NodeID. Leading 1s and first leading 0 of the NodeID
+	// are truncated off. The rest is appended to the IPv6 address (truncated to
+	// 128 bits total).
+	var addr address.Address
+	temp := make([]byte, 0, len(nid))
+	done := false
+	ones := byte(0)
+	bits := byte(0)
+	nBits := 0
+	for idx := 0; idx < 8*len(nid); idx++ {
+		bit := (nid[idx/8] & (0x80 >> byte(idx%8))) >> byte(7-(idx%8))
+		if !done && bit != 0 {
+			ones++
+			continue
+		}
+		if !done && bit == 0 {
+			done = true
+			continue
+		}
+		bits = (bits << 1) | bit
+		nBits++
+		if nBits == 8 {
+			nBits = 0
+			temp = append(temp, bits)
+		}
+	}
+	prefix := address.GetPrefix()
+	copy(addr[:], prefix[:])
+	addr[len(prefix)] = ones
+	copy(addr[len(prefix)+1:], temp)
+	return &addr
+}
diff --git a/main.go b/main.go
index d0519ed..da94066 100644
--- a/main.go
+++ b/main.go
@@ -16,7 +16,6 @@ import (
 	"regexp"
 	"runtime"
 
-	"github.com/yggdrasil-network/yggdrasil-go/src/address"
 	"github.com/yggdrasil-network/yggdrasil-go/src/crypto"
 )
 
@@ -74,7 +73,7 @@ func doBoxKeys(out chan<- keySet) {
 	for {
 		pub, priv := crypto.NewBoxKeys()
 		id := crypto.GetNodeID(pub)
-		ip := net.IP(address.AddrForNodeID(id)[:]).String()
+		ip := net.IP(AddrForNodeID(id)[:]).String()
 		out <- keySet{priv[:], pub[:], id[:], ip}
 	}
 }