Refactor

nixberg · May 11, 2024 · 2f472e5 · 2f472e5
1 parent 5f2f120
commit 2f472e5
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 47 deletions.
diff --git a/Sources/ChaCha8Rand/BinaryEncoding.swift b/Sources/ChaCha8Rand/BinaryEncoding.swift
@@ -13,18 +13,25 @@ extension ChaCha8Rand {
     public init?(decoding buffer: UnsafeRawBufferPointer) {
         precondition(buffer.count == 48, "TODO")
 
-        guard buffer.loadLittleEndianUInt64(fromByteOffset: 00) == Self.header else {
+        guard buffer.loadUnaligned(as: UInt64.self, endianess: .little) == Self.header else {
             return nil
         }
 
-        let (counter, index) = buffer.loadBigEndianUInt64(fromByteOffset: 08)
-            .quotientAndRemainder(dividingBy: 32)
+        let (counter, index) = buffer.loadUnaligned(
+            fromByteOffset: 08,
+            as: UInt64.self,
+            endianess: .big
+        ).quotientAndRemainder(dividingBy: 32)
 
         let seed = ContiguousArray<UInt32>(unsafeUninitializedCapacity: 8) { seed, count in
             seed.initialize(repeating: 0)
             var offset = 16
             for index in seed.indices {
-                seed[index] = buffer.loadLittleEndianUInt32(fromByteOffset: offset)
+                seed[index] = buffer.loadUnaligned(
+                    fromByteOffset: offset,
+                    as: UInt32.self,
+                    endianess: .little
+                )
                 offset &+= 4
             }
             count = 8
@@ -36,13 +43,14 @@ extension ChaCha8Rand {
     public func encode(into buffer: UnsafeMutableRawBufferPointer) {
         precondition(buffer.count == 48, "TODO")
 
-        buffer.storeLittleEndianBytes(of: Self.header, toByteOffset: 00)
+        buffer.storeBytes(of: Self.header, endianess: .little, toByteOffset: 00)
 
-        buffer.storeBigEndianBytes(of: UInt64(counter / 4) * 32 + UInt64(index), toByteOffset: 08)
+        let used = UInt64(counter / 4) * 32 + UInt64(index)
+        buffer.storeBytes(of: used, endianess: .big, toByteOffset: 08)
 
         var offset = 16
         for word in seed {
-            buffer.storeLittleEndianBytes(of: word, toByteOffset: offset)
+            buffer.storeBytes(of: word, endianess: .little, toByteOffset: offset)
             offset &+= 4
         }
     }

diff --git a/Sources/ChaCha8Rand/ChaCha8Rand.swift b/Sources/ChaCha8Rand/ChaCha8Rand.swift
@@ -11,7 +11,8 @@ public struct ChaCha8Rand: RandomNumberGenerator {
         seed = ContiguousArray(unsafeUninitializedCapacity: 8) { seed, count in
             seed.initialize(repeating: 0)
             for index in seed.indices {
-                seed[index] = SystemRandomNumberGenerator.next()
+                var generator = SystemRandomNumberGenerator()
+                seed[index] = generator.next()
             }
             count = 8
         }
@@ -28,14 +29,14 @@ public struct ChaCha8Rand: RandomNumberGenerator {
         precondition(bytes.count == 32, "TODO")
         self.seed = ContiguousArray(unsafeUninitializedCapacity: 8) { seed, count in
             seed.initialize(repeating: 0)
-            seed[0] = bytes.loadLittleEndianUInt32(fromByteOffset: 00)
-            seed[1] = bytes.loadLittleEndianUInt32(fromByteOffset: 04)
-            seed[2] = bytes.loadLittleEndianUInt32(fromByteOffset: 08)
-            seed[3] = bytes.loadLittleEndianUInt32(fromByteOffset: 12)
-            seed[4] = bytes.loadLittleEndianUInt32(fromByteOffset: 16)
-            seed[5] = bytes.loadLittleEndianUInt32(fromByteOffset: 20)
-            seed[6] = bytes.loadLittleEndianUInt32(fromByteOffset: 24)
-            seed[7] = bytes.loadLittleEndianUInt32(fromByteOffset: 28)
+            seed[0] = bytes.loadUnaligned(fromByteOffset: 00, as: UInt32.self, endianess: .little)
+            seed[1] = bytes.loadUnaligned(fromByteOffset: 04, as: UInt32.self, endianess: .little)
+            seed[2] = bytes.loadUnaligned(fromByteOffset: 08, as: UInt32.self, endianess: .little)
+            seed[3] = bytes.loadUnaligned(fromByteOffset: 12, as: UInt32.self, endianess: .little)
+            seed[4] = bytes.loadUnaligned(fromByteOffset: 16, as: UInt32.self, endianess: .little)
+            seed[5] = bytes.loadUnaligned(fromByteOffset: 20, as: UInt32.self, endianess: .little)
+            seed[6] = bytes.loadUnaligned(fromByteOffset: 24, as: UInt32.self, endianess: .little)
+            seed[7] = bytes.loadUnaligned(fromByteOffset: 28, as: UInt32.self, endianess: .little)
             count = 8
         }
         self.block()
@@ -114,10 +115,3 @@ extension UInt64 {
         (UInt32(truncatingIfNeeded: self), UInt32(truncatingIfNeeded: self >> 32))
     }
 }
-
-extension SystemRandomNumberGenerator {
-    fileprivate static func next() -> UInt32 {
-        var generator = Self()
-        return generator.next()
-    }
-}
diff --git a/Sources/ChaCha8Rand/Common.swift b/Sources/ChaCha8Rand/Common.swift
@@ -1,33 +1,43 @@
+enum Endianess {
+    case big
+    case little
+    case native
+}
+
 extension UnsafeRawBufferPointer {
     @inline(__always)
-    func loadBigEndianUInt64(fromByteOffset offset: Int) -> UInt64 {
-        UInt64(bigEndian: self.loadUnaligned(fromByteOffset: offset, as: UInt64.self))
-    }
-
-    @inline(__always)
-    func loadLittleEndianUInt32(fromByteOffset offset: Int) -> UInt32 {
-        UInt32(littleEndian: self.loadUnaligned(fromByteOffset: offset, as: UInt32.self))
-    }
-
-    @inline(__always)
-    func loadLittleEndianUInt64(fromByteOffset offset: Int) -> UInt64 {
-        UInt64(littleEndian: self.loadUnaligned(fromByteOffset: offset, as: UInt64.self))
+    func loadUnaligned<T: FixedWidthInteger>(
+        fromByteOffset offset: Int = 0,
+        as type: T.Type,
+        endianess: Endianess
+    ) -> T {
+        let value = self.loadUnaligned(fromByteOffset: offset, as: type)
+        return switch endianess {
+        case .big:
+            T(bigEndian: value)
+        case .little:
+            T(littleEndian: value)
+        case .native:
+            value
+        }
     }
 }
 
 extension UnsafeMutableRawBufferPointer {
     @inline(__always)
-    func storeBigEndianBytes(of value: UInt64, toByteOffset offset: Int) {
-        self.storeBytes(of: value.bigEndian, toByteOffset: offset, as: UInt64.self)
-    }
-
-    @inline(__always)
-    func storeLittleEndianBytes(of value: UInt32, toByteOffset offset: Int) {
-        self.storeBytes(of: value.littleEndian, toByteOffset: offset, as: UInt32.self)
-    }
-
-    @inline(__always)
-    func storeLittleEndianBytes(of value: UInt64, toByteOffset offset: Int) {
-        self.storeBytes(of: value.littleEndian, toByteOffset: offset, as: UInt64.self)
+    func storeBytes<T: FixedWidthInteger>(
+        of value: T,
+        endianess: Endianess,
+        toByteOffset offset: Int = 0
+    ) {
+        let value = switch endianess {
+        case .big:
+            value.bigEndian
+        case .little:
+            value.littleEndian
+        case .native:
+            value
+        }
+        self.storeBytes(of: value, toByteOffset: offset, as: T.self)
     }
 }