assembly for more string stuff

turbocrime · turbocrime · commit afb34d7bab5e · 2024-11-22T22:55:48.000-08:00
diff --git a/src/CborDecode.sol b/src/CborDecode.sol
@@ -111,7 +111,37 @@ library CBORDecoder {
         return (slice, byteIdx + len);
     }
 
-    function readStringBytes1(bytes memory cborData, uint byteIdx) internal pure returns (bytes1, uint) {
+    function readStringBytes9(bytes memory cborData, uint byteIdx) internal pure returns (bytes9, uint) {
+        uint8 major;
+        uint64 len;
+
+        (major, len, byteIdx) = parseCborHeader(cborData, byteIdx);
+        require(major == MajorText && len <= 9, "expected string length of 9 or less");
+
+        bytes9 slice;
+        assembly {
+            slice := and(mload(add(cborData, add(0x20, byteIdx))), not(shr(mul(len, 8), not(0))))
+        }
+
+        return (slice, byteIdx + len);
+    }
+
+    function readStringBytes7(bytes memory cborData, uint byteIdx) internal pure returns (bytes7, uint) {
+        uint8 major;
+        uint64 len;
+
+        (major, len, byteIdx) = parseCborHeader(cborData, byteIdx);
+        require(major == MajorText && len <= 7, "expected string length of 7 or less");
+
+        bytes7 slice;
+        assembly {
+            slice := and(mload(add(cborData, add(0x20, byteIdx))), not(shr(mul(len, 8), not(0))))
+        }
+
+        return (slice, byteIdx + len);
+    }
+
+    function readStringBytes1_normal(bytes memory cborData, uint byteIdx) internal pure returns (bytes1, uint) {
         uint8 major;
         uint64 len;
 
@@ -122,6 +152,24 @@ library CBORDecoder {
         return (cborData[byteIdx], byteIdx + len);
     }
 
+    function readStringBytes1(bytes memory cborData, uint byteIdx) internal pure returns (bytes1, uint) {
+        return readStringBytes1_normal(cborData, byteIdx);
+    }
+
+    function readStringBytes1_assembly(bytes memory cborData, uint byteIdx) internal pure returns (bytes1, uint) {
+        uint8 head;
+        assembly {
+            head := mload(add(cborData, add(0x20, byteIdx)))
+        }
+        require(head == MajorText << shiftMajor | 0x01, "expected string length of 1");
+        bytes1 slice;
+        byteIdx += 1;
+        assembly {
+            slice := mload(add(cborData, add(0x20, byteIdx)))
+        }
+        return (slice, byteIdx + 2);
+    }
+
     function skipString(bytes memory cborData, uint byteIdx) internal pure returns (uint) {
         uint8 major;
         uint64 len;
diff --git a/src/CommitCbor.sol b/src/CommitCbor.sol
@@ -23,9 +23,9 @@ library CommitCbor {
         require(mapLen == 5, "expected 5 fields in commit");
 
         for (uint i = 0; i < mapLen; i++) {
-            bytes memory mapKey;
-            (mapKey, byteIdx) = cborData.readStringBytes(byteIdx);
-            if (bytes8(mapKey) == "version") {
+            bytes7 mapKey;
+            (mapKey, byteIdx) = cborData.readStringBytes7(byteIdx);
+            if (bytes7(mapKey) == "version") {
                 (ret.version, byteIdx) = cborData.readUInt8(byteIdx);
                 require(ret.version == COMMIT_VERSION, "commit version number must be 3");
             } else if (bytes5(mapKey) == "data") {
diff --git a/src/RecordCbor.sol b/src/RecordCbor.sol
@@ -18,8 +18,8 @@ library RecordCbor {
         require(mapLen == 4, "expected 4 fields in record");
 
         for (uint i = 0; i < mapLen; i++) {
-            bytes memory mapKey;
-            (mapKey, byteIdx) = cborData.readStringBytes(byteIdx);
+            bytes9 mapKey;
+            (mapKey, byteIdx) = cborData.readStringBytes9(byteIdx);
             if (bytes5(mapKey) == "text") {
                 (ret.text, byteIdx) = cborData.readString(byteIdx);
             } else if (bytes6(mapKey) == "$type") {
@@ -31,7 +31,7 @@ library RecordCbor {
                     // langs string unused.
                     byteIdx = cborData.skipString(byteIdx);
                 }
-            } else if (bytes10(mapKey) == "createdAt") {
+            } else if (bytes9(mapKey) == "createdAt") {
                 // createdAt string unused.
                 // this field is arbitrary user-defined data. the useful and
                 // verifiable timestamp is the signed commit repo revision field
diff --git a/src/TreeNodeCbor.sol b/src/TreeNodeCbor.sol
@@ -65,6 +65,40 @@ library TreeNodeCbor {
     }
 
     function buildEntryKeys(TreeNodeE[] memory e) internal pure returns (TreeNodeEntry[] memory) {
+        return buildEntryKeys_assembly(e);
+    }
+
+    function buildEntryKeys_assembly(TreeNodeE[] memory e) internal pure returns (TreeNodeEntry[] memory) {
+        TreeNodeEntry[] memory entries = new TreeNodeEntry[](e.length);
+        bytes memory previousKey = new bytes(0);
+        for (uint i = 0; i < e.length; i++) {
+            uint8 p = e[i].p;
+            bytes memory k = e[i].k;
+            bytes memory key = new bytes(p + k.length);
+            //console.log("i=%s", i);
+            //console.log("before assembly p=%s k.length=%s key.length=%s", p, k.length, key.length);
+            // Calculate number of words needed
+            uint pWords = (p + 31) / 32; // ceil(p/32)
+            uint kWords = (k.length + 31) / 32; // ceil(k.length/32)
+            uint j;
+
+            assembly {
+                for { j := 0 } lt(j, pWords) { j := add(j, 1) } {
+                    mstore(add(key, add(0x20, mul(j, 32))), mload(add(previousKey, add(0x20, mul(j, 32)))))
+                }
+
+                for { j := 0 } lt(j, kWords) { j := add(j, 1) } {
+                    mstore(add(add(key, 0x20), add(p, mul(j, 32))), mload(add(k, add(0x20, mul(j, 32)))))
+                }
+            }
+            //console.log("after assembly key=%s key.length=%s", string(key), key.length);
+            entries[i] = TreeNodeEntry(string(key), e[i].v, e[i].t);
+            previousKey = key;
+        }
+        return entries;
+    }
+
+    function buildEntryKeys_loop(TreeNodeE[] memory e) internal pure returns (TreeNodeEntry[] memory) {
         TreeNodeEntry[] memory entries = new TreeNodeEntry[](e.length);
         bytes memory previousKey = new bytes(0);
         for (uint i = 0; i < e.length; i++) {
diff --git a/test/CborDecode.t.sol b/test/CborDecode.t.sol
@@ -169,19 +169,19 @@ contract CborDecodeTest {
         bytes1 mapKey;
         uint8 mapValue;
         (mapKey, index) = input.readStringBytes1(index);
-        require(mapKey == bytes1("a"), "map key is not 'a'");
+        require(mapKey == "a", "map key is not 'a'");
 
         (mapValue, index) = input.readUInt8(index);
         require(mapValue == 1, "map value is not 1");
 
         (mapKey, index) = input.readStringBytes1(index);
-        require(mapKey == bytes1("b"), "map key is not 'b'");
+        require(mapKey == "b", "map key is not 'b'");
 
         (mapValue, index) = input.readUInt8(index);
         require(mapValue == 2, "map value is not 2");
 
         (mapKey, index) = input.readStringBytes1(index);
-        require(mapKey == bytes1("c"), "map key is not 'c'");
+        require(mapKey == "c", "map key is not 'c'");
 
         (mapValue, index) = input.readUInt8(index);
         require(mapValue == 3, "map value is not 3");
