Skip to content

fix(stakeweight): correct permanent lock epoch accounting #45

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rplusq merged 2 commits into from
May 12, 2026
Merged

fix(stakeweight): correct permanent lock epoch accounting #45

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
615f1a0
fix(stakeweight): correct permanent lock epoch accounting and add reg…
rplusq Mar 2, 2026
343079b
chore(stakeweight): drop archived buggy contract and demo test
rplusq May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 64 additions & 81 deletions evm/src/StakeWeight.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,13 +113,13 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
mapping(address => bool) isPermanent; // current permanent flag
mapping(address => uint256) permanentBaseWeeks; // chosen discrete weeks
mapping(address => uint256) permanentStakeWeight; // current user's permanent stake weight (amount * baseWeeks /
// maxWeeks)
// maxWeeks)
// Global accumulator (stake weight units, not tokens)
uint256 permanentTotalSupply; // sum of all permanent stake weight
// Parallel histories (align with existing epochs / userEpochs)
mapping(uint256 => uint256) globalPermanentSupplyAtEpoch; // epoch -> permanentTotalSupply at that global point
mapping(address => mapping(uint256 => uint256)) userPermanentWeightAtEpoch; // user, userEpoch -> permanent
// stake weight at that user point
// stake weight at that user point
}

function _getStakeWeightStorage() internal pure returns (StakeWeightStorage storage s) {
Expand Down Expand Up @@ -536,8 +536,8 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
// As we cannot figure that out block timestamp -> block number exactly
// when query states from xxxAt methods, we need to calculate block number
// based on initalLastPoint
lastPoint.blockNumber =
initialLastPoint.blockNumber + ((blockSlope * ((weekCursor - initialLastPoint.timestamp))) / MULTIPLIER);
lastPoint.blockNumber = initialLastPoint.blockNumber
+ ((blockSlope * ((weekCursor - initialLastPoint.timestamp))) / MULTIPLIER);
epoch_ = epoch_ + 1;
if (weekCursor == block.timestamp) {
// Hard to be happened, but better handling this case too
Expand Down Expand Up @@ -655,9 +655,7 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
revert TransferRestrictionsEnabled();
}
LockedBalance memory lock = LockedBalance({
amount: s.locks[for_].amount,
end: s.locks[for_].end,
transferredAmount: s.locks[for_].transferredAmount
amount: s.locks[for_].amount, end: s.locks[for_].end, transferredAmount: s.locks[for_].transferredAmount
});

if (for_ == address(0)) revert InvalidAddress(for_);
Expand Down Expand Up @@ -692,9 +690,7 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard

// Store prevLocked
LockedBalance memory newLocked = LockedBalance({
amount: prevLocked.amount,
end: prevLocked.end,
transferredAmount: prevLocked.transferredAmount
amount: prevLocked.amount, end: prevLocked.end, transferredAmount: prevLocked.transferredAmount
});

// Adding new lock to existing lock, or if lock is expired
Expand Down Expand Up @@ -816,6 +812,7 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
function increaseLockAmount(uint256 amount) external nonReentrant {
StakeWeightStorage storage s = _getStakeWeightStorage();
_requireNotPaused(s);
if (s.isPermanent[msg.sender]) revert AlreadyPermanent();
_increaseLockAmount(msg.sender, amount, true);
}

Expand All @@ -837,7 +834,7 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
LockedBalance memory lock = s.locks[for_];
if (amount == 0) revert InvalidAmount(amount);
if (lock.amount == 0) revert NonExistentLock();
if (lock.end <= block.timestamp) revert ExpiredLock(block.timestamp, lock.end);
if (!s.isPermanent[for_] && lock.end <= block.timestamp) revert ExpiredLock(block.timestamp, lock.end);
_depositFor(for_, amount, 0, lock, ACTION_INCREASE_LOCK_AMOUNT, isTransferred);
}

Expand Down Expand Up @@ -904,6 +901,17 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
// Store base weeks for later unlock
s.permanentBaseWeeks[msg.sender] = baseWeeks;

// Calculate permanent stake weight using same formula as bias calculation
// permanentWeight = amount * duration / MAX_LOCK_CAP
// This ensures consistency with decaying positions
uint256 amount = SafeCast.toUint256(lock.amount);
uint256 permanentWeight = Math.mulDiv(amount, duration, MAX_LOCK_CAP);

// Update permanent state BEFORE checkpoints so intermediate epochs have correct supply
s.isPermanent[msg.sender] = true;
s.permanentStakeWeight[msg.sender] = permanentWeight;
s.permanentTotalSupply += permanentWeight;

// TWO-PHASE CHECKPOINT for clean conversion:
// Phase 1: Remove all decaying weight (amount -> 0)
LockedBalance memory zeroLock = LockedBalance({
Expand All @@ -924,21 +932,8 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
// CRITICAL: Persist the permanent lock with amount restored
s.locks[msg.sender] = permanentLock;

// Calculate permanent stake weight using same formula as bias calculation
// permanentWeight = amount * duration / MAX_LOCK_CAP
// This ensures consistency with decaying positions
uint256 amount = SafeCast.toUint256(lock.amount);
uint256 permanentWeight = Math.mulDiv(amount, duration, MAX_LOCK_CAP);

// Update permanent state
s.isPermanent[msg.sender] = true;
s.permanentStakeWeight[msg.sender] = permanentWeight;
s.permanentTotalSupply += permanentWeight;

// Record in parallel histories (after checkpoint incremented epoch)
uint256 userEpoch = s.userPointEpoch[msg.sender];
s.userPermanentWeightAtEpoch[msg.sender][userEpoch] = permanentWeight;
// Use pointHistory.length - 1 as the checkpoint was just pushed
// Record in parallel histories AFTER checkpoints created new epochs
s.userPermanentWeightAtEpoch[msg.sender][s.userPointEpoch[msg.sender]] = permanentWeight;
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;

emit PermanentConversion(msg.sender, duration, block.timestamp);
Expand All @@ -962,23 +957,21 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
LockedBalance memory next =
LockedBalance({ amount: prev.amount, end: newEnd, transferredAmount: prev.transferredAmount });

// Run checkpoint to re-introduce slope and bias
_checkpoint(msg.sender, prev, next);

// CRITICAL: Persist the lock change with new end time
s.locks[msg.sender] = next;

// Remove permanent stake weight
// Remove permanent stake weight BEFORE checkpoint so intermediate epochs have correct supply
uint256 permanentWeight = s.permanentStakeWeight[msg.sender];
s.permanentTotalSupply -= permanentWeight;
s.isPermanent[msg.sender] = false;
s.permanentStakeWeight[msg.sender] = 0;
s.permanentBaseWeeks[msg.sender] = 0; // Clear the duration

// Update parallel histories
uint256 userEpoch = s.userPointEpoch[msg.sender];
s.userPermanentWeightAtEpoch[msg.sender][userEpoch] = 0;
// Use pointHistory.length - 1 as the checkpoint was just pushed
// Run checkpoint to re-introduce slope and bias
_checkpoint(msg.sender, prev, next);

// CRITICAL: Persist the lock change with new end time
s.locks[msg.sender] = next;

// Update parallel histories AFTER checkpoint created new epochs
s.userPermanentWeightAtEpoch[msg.sender][s.userPointEpoch[msg.sender]] = 0;
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;

emit UnlockTriggered(msg.sender, newEnd, block.timestamp);
Expand Down Expand Up @@ -1348,8 +1341,13 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
uint256 currentBaseWeeks = s.permanentBaseWeeks[msg.sender];
uint256 newBaseWeeks = _validatePermanentDuration(newDuration);

if (newBaseWeeks != currentBaseWeeks && newBaseWeeks < currentBaseWeeks) {
revert InvalidDuration(newDuration); // Simplified: duration must increase when changing
// Must add tokens or increase duration - otherwise no-op
if (amount == 0 && newBaseWeeks == currentBaseWeeks) {
revert InvalidAmount(0);
}

if (newBaseWeeks < currentBaseWeeks) {
revert InvalidDuration(newDuration);
}

// Store previous lock state for checkpointing
Expand All @@ -1371,46 +1369,35 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
lock.transferredAmount += amount;
}

// Update permanent weight if amount added or duration changed
if (amount > 0 || newBaseWeeks != currentBaseWeeks) {
// Calculate old weight to remove
uint256 oldWeight = s.permanentStakeWeight[msg.sender];

// Calculate new weight based on total amount and new duration
uint256 totalAmount = SafeCast.toUint256(int256(lock.amount));
uint256 newWeight = Math.mulDiv(totalAmount, newBaseWeeks * 1 weeks, MAX_LOCK_CAP);

// Update user's permanent weight
s.permanentStakeWeight[msg.sender] = newWeight;
// Calculate old weight to remove
uint256 oldWeight = s.permanentStakeWeight[msg.sender];

// Update global permanent supply
s.permanentTotalSupply = s.permanentTotalSupply - oldWeight + newWeight;
// Calculate new weight based on total amount and new duration
uint256 totalAmount = SafeCast.toUint256(int256(lock.amount));
uint256 newWeight = Math.mulDiv(totalAmount, newBaseWeeks * 1 weeks, MAX_LOCK_CAP);

// Record history at current global epoch
uint256 currentEpoch = s.epoch;
if (currentEpoch > 0) {
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;
}
// Update user's permanent weight
s.permanentStakeWeight[msg.sender] = newWeight;

// Record user permanent history at their current epoch
uint256 userEpoch = s.userPointEpoch[msg.sender];
if (userEpoch > 0) {
s.userPermanentWeightAtEpoch[msg.sender][userEpoch] = newWeight;
}
// Update global permanent supply
s.permanentTotalSupply = s.permanentTotalSupply - oldWeight + newWeight;

// Update duration if changed
if (newBaseWeeks != currentBaseWeeks) {
s.permanentBaseWeeks[msg.sender] = newBaseWeeks;
emit DurationIncreased(msg.sender, newDuration, block.timestamp);
}
// Update duration if changed
if (newBaseWeeks != currentBaseWeeks) {
s.permanentBaseWeeks[msg.sender] = newBaseWeeks;
emit DurationIncreased(msg.sender, newDuration, block.timestamp);
}

// Save updated lock
s.locks[msg.sender] = lock;

// Checkpoint the change
// Checkpoint FIRST to create new epoch
_checkpoint(msg.sender, prevLock, lock);

// THEN record histories at the NEW epochs (after checkpoint created them)
s.userPermanentWeightAtEpoch[msg.sender][s.userPointEpoch[msg.sender]] = newWeight;
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;

// Emit deposit event if amount was added
if (amount > 0) {
emit Deposit(msg.sender, amount, 0, ACTION_UPDATE_LOCK, amount, block.timestamp);
Expand Down Expand Up @@ -1450,28 +1437,24 @@ contract StakeWeight is Initializable, AccessControlUpgradeable, ReentrancyGuard
// Update global permanent supply
s.permanentTotalSupply = s.permanentTotalSupply - oldWeight + newWeight;

// Record history at current global epoch
uint256 currentEpoch = s.epoch;
if (currentEpoch > 0) {
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;
}

// Record user permanent history at their current epoch
uint256 userEpoch = s.userPointEpoch[msg.sender];
if (userEpoch > 0) {
s.userPermanentWeightAtEpoch[msg.sender][userEpoch] = newWeight;
}

// Store previous lock state for checkpointing (no amount changes)
LockedBalance memory prevLock =
LockedBalance({ amount: lock.amount, end: lock.end, transferredAmount: lock.transferredAmount });

// Checkpoint the change (lock doesn't change, but we need to update global state)
// FIX: Checkpoint FIRST to create new epoch
_checkpoint(msg.sender, prevLock, lock);

// FIX: THEN record histories at the NEW epochs (after checkpoint created them)
s.userPermanentWeightAtEpoch[msg.sender][s.userPointEpoch[msg.sender]] = newWeight;
s.globalPermanentSupplyAtEpoch[s.pointHistory.length - 1] = s.permanentTotalSupply;

emit DurationIncreased(msg.sender, newDuration, block.timestamp);
}

/*//////////////////////////////////////////////////////////////////////////
VIEW FUNCTIONS
//////////////////////////////////////////////////////////////////////////*/

/// @notice Get the total permanent supply
/// @return The total permanent supply across all users
function permanentSupply() external view returns (uint256) {
Expand Down
67 changes: 67 additions & 0 deletions evm/src/StakeWeightHealer.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
// SPDX-License-Identifier: MIT
pragma solidity 0.8.25;

import { AccessControlUpgradeable } from "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
import { StakeWeight } from "./StakeWeight.sol";

/**
* @title StakeWeightHealer
* @notice Minimal contract for healing affected permanent lock users
*
* @dev This contract is used in a "sandwich upgrade" pattern:
* 1. Upgrade proxy to StakeWeightHealer
* 2. Admin calls batchHealPermanentWeights() to fix affected users
* 3. Upgrade proxy back to StakeWeight
*
* IMPORTANT: This contract accesses the SAME storage slot as StakeWeight
* via ERC-7201 namespaced storage, reusing StakeWeight.StakeWeightStorage
* directly. It does NOT inherit from StakeWeight to stay minimal.
*
* @author WalletConnect
*/
contract StakeWeightHealer is AccessControlUpgradeable {
/*//////////////////////////////////////////////////////////////////////////
STORAGE
//////////////////////////////////////////////////////////////////////////*/

bytes32 private constant STAKE_WEIGHT_STORAGE_POSITION = keccak256("com.walletconnect.stakeweight.storage");

function _getStakeWeightStorage() internal pure returns (StakeWeight.StakeWeightStorage storage s) {
bytes32 position = STAKE_WEIGHT_STORAGE_POSITION;
assembly {
s.slot := position
}
}

/*//////////////////////////////////////////////////////////////////////////
EVENTS
//////////////////////////////////////////////////////////////////////////*/

event PermanentWeightHealed(address indexed user, uint256 epoch, uint256 weight);

/*//////////////////////////////////////////////////////////////////////////
HEALING FUNCTIONS
//////////////////////////////////////////////////////////////////////////*/

/**
* @notice Allows admin to batch heal multiple affected users
* @param users Array of user addresses to heal
*/
function batchHealPermanentWeights(address[] calldata users) external onlyRole(DEFAULT_ADMIN_ROLE) {
StakeWeight.StakeWeightStorage storage s = _getStakeWeightStorage();

for (uint256 i = 0; i < users.length; i++) {
address user = users[i];

if (!s.isPermanent[user]) continue;

uint256 currentWeight = s.permanentStakeWeight[user];
uint256 userEpoch = s.userPointEpoch[user];

if (currentWeight > 0 && userEpoch > 0 && s.userPermanentWeightAtEpoch[user][userEpoch] == 0) {
s.userPermanentWeightAtEpoch[user][userEpoch] = currentWeight;
emit PermanentWeightHealed(user, userEpoch, currentWeight);
}
}
}
}
Loading
Loading