bug in balanceOf somewhere

LICENSE

@@ -1,24 +1,21 @@
-This is free and unencumbered software released into the public domain.
+MIT License
 
-Anyone is free to copy, modify, publish, use, compile, sell, or
-distribute this software, either in source code form or as a compiled
-binary, for any purpose, commercial or non-commercial, and by any
-means.
+Copyright (c) 2023 Erhan Tezcan
 
-In jurisdictions that recognize copyright laws, the author or authors
-of this software dedicate any and all copyright interest in the
-software to the public domain. We make this dedication for the benefit
-of the public at large and to the detriment of our heirs and
-successors. We intend this dedication to be an overt act of
-relinquishment in perpetuity of all present and future rights to this
-software under copyright law.
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
-IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
-OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
-ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
-OTHER DEALINGS IN THE SOFTWARE.
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
 
-For more information, please refer to <http://unlicense.org/>
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -8,6 +8,7 @@ The concept works for any order actually, but we would like to use an order that
 
 - have polynomial for approvals too?
 - maybe use FFI to generate the basis based on total supply via script?
+- bug with memory offsets somehow?
 
 ## Usage
 

src/Huffd1.huff

@@ -9,20 +9,24 @@
 #include "./util/Polynomial.huff"
 
 ///////////////////////////////////////////////////////////////////////////////
-////                             INTERFACES                                ////
+////                                TABLE                                  ////
 ///////////////////////////////////////////////////////////////////////////////
-#define function name() nonpayable returns (string)
-#define function symbol() nonpayable returns (string)
-#define function ownerOf(uint256) view returns (address)
-
-/// Table
 #define table Basis {
     0x01050700020c000607
 }
 #define constant TOTAL_SUPPLY = 0x03
 #define constant ORDER = 0x0d
 #define constant COEFF_SIZE = 0x01
 
+///////////////////////////////////////////////////////////////////////////////
+////                             INTERFACES                                ////
+///////////////////////////////////////////////////////////////////////////////
+#define function name() nonpayable returns (string)
+#define function symbol() nonpayable returns (string)
+#define function ownerOf(uint256) view returns (address)
+#define function transfer(address, uint256) nonpayable returns ()
+#define function balanceOf(address) view returns (uint256)
+
 ///////////////////////////////////////////////////////////////////////////////
 ////                               STORAGE                                 ////
 ///////////////////////////////////////////////////////////////////////////////
@@ -60,6 +64,89 @@
     // output:             [a] (owner)
 }
 
+#define macro TRANSFER() = takes (3) returns (0) {
+    // input:           // [t, b, a] (token, to, from)
+
+    // assert that token is owner by caller
+    dup3 OWNER_OF()     // [t, b, a, o] (owner of t)
+    dup2 eq             // [t, b, a, a == o]
+    success jumpi
+    0x00 0x00 revert
+    success:
+
+    // find scale amount
+    [ORDER] sub         // [t, b, -a] (-a = p - a) in field
+    add                 // [t, v] (v := b - a)
+
+    // load balance polynomial
+    [BAL_SLOT]          // [t, v, s] (storageSlot)
+    0x00                // [t, v, s, mP] (memoryOffset)
+    [TOTAL_SUPPLY]      // [t, v, s, mP, l] (length)
+    POLY_SLOAD()        // [t, v, mP, l]
+
+    // load basis polynomial
+    __tablestart(Basis) // [t, v, mP, l, c]
+    dup5                // [t, v, mP, l, c, t]
+    POLY_OFFSET(COEFF_SIZE, TOTAL_SUPPLY)
+                        // [t, v, mP, l, c, o]
+    add                 // [t, v, mP, l, c] (codeOffset)
+    [COEFF_SIZE]        // [t, v, mP, l, c, s] (coefficient size)
+    0xFF                // [t, v, mP, l, c, s, mQ] (memOffset)
+    dup4                // [t, v, mP, l, c, s, mQ, l] (length)
+    POLY_CODECOPY()     // [t, v, mP, l, mQ, l]
+    pop swap1           // [t, v, mP, mQ, l]
+
+    // scale basis polynomial with (b - a)
+    dup4                // [t, v, mP, mQ, l, v]
+    POLY_SCALE(ORDER)   // [t, v, mP, mQ, l]
+
+    // add scaled basis polynomial to balance polynomial
+    POLY_ADD(ORDER)     // [t, v, mP, l]
+
+    // store the new polynomial
+    [BAL_SLOT] swap1    // [t, v, mP, sQ, l]
+    POLY_SSTORE()       // [t, v, sQ, l]
+    pop pop pop pop     // []
+}
+
+#define macro BALANCE_OF() = takes (1) returns (1) {
+    // input:           // [a] (address)
+    0x00                // [a, s] (sum; total balance initially 0)
+    [TOTAL_SUPPLY]      // [a, s, l] (length; number of tokens)
+    0x00                // [a, s, l, i] (index starts with 0)
+
+    // TODO: can load the poly only once at the start
+    loop_begin:
+    // [a, s, l, i]
+
+        // loop condition
+        dup2 dup2       // [a, s, l, i, l, i]
+        lt iszero       // [a, s, l, i, i >= l]
+        loop_end jumpi 
+
+        // check ownership for token i
+        dup1            // [a, s, l, i, i] (i as token id)
+        OWNER_OF()      // [a, s, l, i, o] (owner of i)
+        dup5            // [a, s, l, i, o, a]
+        eq              // [a, s, l, i, eq] (eq := o == a)  
+
+        // add to total balance
+        dup4            // [a, s,  l, i, eq, s]
+        add             // [a, s,  l, i, s'] (s' := s + eq)
+        swap3           // [a, s', l, i, s] 
+        pop             // [a, s,  l, i]
+
+        // increment index
+        0x01 add        // [a, s', l, i'] (i' := i + 1)
+
+    // [a, s, l, i]
+    loop_begin jump
+    loop_end:
+
+    pop pop swap1 pop   // [s]
+    // output:          // [s] (total balance)
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 ////                               MAIN                                    ////
 ///////////////////////////////////////////////////////////////////////////////
@@ -70,10 +157,12 @@
     [BAL_SLOT] sstore   // []
 }
 
-#define macro HUFFD1_MAIN() = takes (1) returns (0) {
+#define macro HUFFD1_MAIN() = takes (1) returns (1) {
     dup1 __FUNC_SIG(name)      eq name      jumpi
     dup1 __FUNC_SIG(symbol)    eq symbol    jumpi
     dup1 __FUNC_SIG(ownerOf)   eq ownerOf   jumpi
+    dup1 __FUNC_SIG(transfer)  eq transfer  jumpi
+    dup1 __FUNC_SIG(balanceOf) eq balanceOf jumpi
     no_match jump
 
     name:
@@ -85,6 +174,17 @@
         OWNER_OF()          // [a] (address)
         0x00 mstore         // []
         0x20 0x00 return    // []
+    transfer:
+        0x24 calldataload   // [t] (tokenId)
+        0x04 calldataload   // [t, b] (to address)
+        caller              // [t, b, a] (from address)
+        TRANSFER()          // []
+        0x00 0x00 return    // []
+    balanceOf:
+        0x04 calldataload   // [a] (address)
+        BALANCE_OF()        // [s] (balance)
+        0x00 mstore         // []   
+        0x20 0x00 return    // []
 
     no_match:
 }

test/Huffd1.sage

@@ -104,7 +104,6 @@ class HUFFD1:
         code = code + "\n#define constant TOTAL_SUPPLY = " + hex(self.total_supply)
         code = code + "\n#define constant COEFF_SIZE = " + hex(self.coeff_size)
         code = code + "\n#define constant ORDER = " + hex(self.order)
-        code = code + "\n"
         path = "./src/data/Basis" + str(self.total_supply) + ".huff"
         with open(path, "w") as _file:
             _file.write(code)
@@ -136,7 +135,10 @@ if __name__ == "__main__":
     # print(huffd1.bs[0])
 
     # test transfer
-    # TODO
+    NEW_OWNER = 0x9
+    assert huffd1.ownerOf(0) == OWNER
+    huffd1.transferFrom(OWNER, NEW_OWNER, 0)
+    assert huffd1.ownerOf(0) == NEW_OWNER
 
     # export table for Huff
     huffd1.export_basis_table()

test/Huffd1.t.sol

@@ -37,6 +37,27 @@ contract Huffd1Test is Test {
         assertEq(OWNER, huffd1.ownerOf(0));
         assertEq(OWNER, huffd1.ownerOf(1));
     }
+
+    /// @dev Should transfer correctly.
+    function test_Transfer() public {
+        assertEq(OWNER, huffd1.ownerOf(0));
+
+        address NEW_OWNER = address(0x9);
+        vm.prank(OWNER);
+        huffd1.transfer(NEW_OWNER, 0);
+        assertEq(huffd1.ownerOf(0), NEW_OWNER);
+    }
+
+    /// @dev Should fail to transfer a not-owned token.
+    function testFail_Transfer() public {
+        vm.prank(address(0x5)); // not an owner
+        huffd1.transfer(OWNER, 0);
+    }
+
+    /// @dev Should give corect balance.
+    function test_BalanceOf() public {
+        assertEq(TOTAL_SUPPLY, huffd1.balanceOf(OWNER));
+    }
 }
 
 // util/Owned.huff

test/Polynomial.t.sol

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: Unlicense
 pragma solidity ^0.8.15;
 
-import "foundry-huff/HuffDeployer.sol";
+import {HuffDeployer} from "foundry-huff/HuffDeployer.sol";
 import {Test} from "forge-std/Test.sol";
-import "forge-std/console.sol";
+import {console} from "forge-std/console.sol";
 
 contract PolynomialTest is Test {
     IPolynomial polynomial;
@@ -27,11 +27,20 @@ contract PolynomialTest is Test {
         }
     }
 
-    // TODO: do more exhaustive tests
-    function test_AddEval() public view {
-        console.log("L0(0) + L0(0): ", polynomial.addEval(0, 0, 0));
-        console.log("L0(0) + L1(0): ", polynomial.addEval(0, 1, 0));
-        console.log("L0(0) + L2(0): ", polynomial.addEval(0, 2, 0));
+    function test_AddEval() public {
+        for (uint256 basis1 = 0; basis1 < TOTAL_SUPPLY; ++basis1) {
+            for (uint256 basis2 = 0; basis2 < TOTAL_SUPPLY; ++basis2) {
+                for (uint256 tokenId = 0; tokenId < TOTAL_SUPPLY; ++tokenId) {
+                    uint256 eval1 = polynomial.eval(basis1, tokenId);
+                    uint256 eval2 = polynomial.eval(basis2, tokenId);
+                    uint256 addEval = polynomial.addEval(basis1, basis2, tokenId);
+                    // console.log("BASIS1:", basis1, "\tBASIS2:", basis2);
+                    // console.log("EVAL1:", eval1, "\tEVAL2:", eval2);
+                    // console.log("TOKEN:", tokenId, "\tADD:", addEval);
+                    assertEq((eval1 + eval2) % ORDER, addEval);
+                }
+            }
+        }
     }
 
     function test_ScaleEval() public {
@@ -48,18 +57,36 @@ contract PolynomialTest is Test {
     }
 
     function test_AddScaleStore() public {
+        uint256 basisToAdd = 0;
+        polynomial.addStore(basisToAdd);
+
+        uint256 scale = 4;
+        polynomial.scaleStore(scale);
+
+        for (uint256 basis = 0; basis < TOTAL_SUPPLY; ++basis) {
+            VERBOSE ? console.log("") : ();
+            for (uint256 tokenId = 0; tokenId < TOTAL_SUPPLY; ++tokenId) {
+                uint256 expected = 0;
+                if (basisToAdd == tokenId) {
+                    expected++;
+                }
+                if (basis == tokenId) {
+                    expected++;
+                }
+                expected = scale * expected;
+                expected = expected % ORDER;
+                assertEq(polynomial.evalStore(tokenId), expected);
+                // VERBOSE ? console.log("BASIS(0) *", scale, "=", scaledEval) : ();
+            }
+        }
         console.log("P(0)", polynomial.evalStore(0));
         console.log("P(1)", polynomial.evalStore(1));
         console.log("P(2)", polynomial.evalStore(2));
 
-        polynomial.addStore(1);
-
         console.log("P(0) + L1(0)", polynomial.evalStore(0));
         console.log("P(1) + L1(1)", polynomial.evalStore(1));
         console.log("P(2) + L1(2)", polynomial.evalStore(2));
 
-        polynomial.scaleStore(5);
-
         console.log("5(P(0) + L1(0))", polynomial.evalStore(0));
         console.log("5(P(1) + L1(1))", polynomial.evalStore(1));
         console.log("5(P(2) + L1(2))", polynomial.evalStore(2));

test/mocks/PolynomialWrappers.huff

@@ -68,7 +68,7 @@
                             // [mP, cc, bb']        (basis offset)
         add                 // [mP, cc']            (cc' := cc + bb)
         [COEFF_SIZE]        // [mP, cc, s]          (coefficient size)
-        0xF0                // [mP, cc, s, mQ]      (memory offset) 
+        0xFFFF              // [mP, cc, s, mQ]      (memory offset) 
         [TOTAL_SUPPLY]      // [mP, cc, s, mQ, l]   (length)
         POLY_CODECOPY()     // [mP, mQ, l]          (mem[m..m+l] := polynomial)