[CT-950] safety heap methods

protocol/x/subaccounts/keeper/safety_heap.go

@@ -0,0 +1,174 @@
+package keeper
+
+import (
+	"cosmossdk.io/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+)
+
+// RemoveSubaccountFromSafetyHeap removes a subaccount from the safety heap
+// given a peretual and side.
+func (k Keeper) RemoveSubaccountFromSafetyHeap(
+	ctx sdk.Context,
+	subaccountId types.SubaccountId,
+	perpetualId uint32,
+	side types.SafetyHeapPositionSide,
+) {
+	store := k.GetSafetyHeapStore(ctx, perpetualId, side)
+	index := k.MustGetSubaccountHeapIndex(store, subaccountId)
+	k.MustRemoveElementAtIndex(ctx, store, index)
+}
+
+// AddSubaccountToSafetyHeap adds a subaccount to the safety heap
+// given a perpetual and side.
+func (k Keeper) AddSubaccountToSafetyHeap(
+	ctx sdk.Context,
+	subaccountId types.SubaccountId,
+	perpetualId uint32,
+	side types.SafetyHeapPositionSide,
+) {
+	store := k.GetSafetyHeapStore(ctx, perpetualId, side)
+	k.Insert(ctx, store, subaccountId)
+}
+
+// Heap methods
+
+// Insert inserts a subaccount into the safety heap.
+func (k Keeper) Insert(
+	ctx sdk.Context,
+	store prefix.Store,
+	subaccountId types.SubaccountId,
+) {
+	// Add the subaccount to the end of the heap.
+	length := k.GetSafetyHeapLength(store)
+	k.SetSubaccountAtIndex(store, subaccountId, length)
+
+	// Increment the size of the heap.
+	k.SetSafetyHeapLength(store, length+1)
+
+	// Heapify up the element at the end of the heap
+	// to restore the heap property.
+	k.HeapifyUp(ctx, store, length)
+}
+
+// MustRemoveElementAtIndex removes the element at the given index
+// from the safety heap.
+func (k Keeper) MustRemoveElementAtIndex(
+	ctx sdk.Context,
+	store prefix.Store,
+	index uint32,
+) {
+	length := k.GetSafetyHeapLength(store)
+	if index >= length {
+		panic(types.ErrSafetyHeapSubaccountNotFoundAtIndex)
+	}
+
+	// Swap the element with the last element.
+	k.Swap(store, index, length-1)
+
+	// Remove the last element.
+	k.DeleteSubaccountAtIndex(store, length-1)
+	k.SetSafetyHeapLength(store, length-1)
+
+	// Heapify down the element at the given index
+	// to restore the heap property.
+	k.HeapifyDown(ctx, store, index)
+}
+
+// HeapifyUp moves the element at the given index up the heap
+// until the heap property is restored.
+func (k Keeper) HeapifyUp(
+	ctx sdk.Context,
+	store prefix.Store,
+	index uint32,
+) {
+	if index == 0 {
+		return
+	}
+
+	parentIndex := (index - 1) / 2
+	if k.Less(ctx, store, index, parentIndex) {
+		k.Swap(store, index, parentIndex)
+		k.HeapifyUp(ctx, store, parentIndex)
+	}
+}
+
+// HeapifyDown moves the element at the given index down the heap
+// until the heap property is restored.
+func (k Keeper) HeapifyDown(
+	ctx sdk.Context,
+	store prefix.Store,
+	index uint32,
+) {
+	leftIndex, rightIndex := 2*index+1, 2*index+2
+
+	length := k.GetSafetyHeapLength(store)
+	if rightIndex < length && k.Less(ctx, store, rightIndex, leftIndex) {
+		// Compare the current node with the right child
+		// if right child exists and is less than the left child.
+		if k.Less(ctx, store, rightIndex, index) {
+			k.Swap(store, index, rightIndex)
+			k.HeapifyDown(ctx, store, rightIndex)
+		}
+	} else if leftIndex < length {
+		// Compare the current node with the left child
+		// if left child exists.
+		if k.Less(ctx, store, leftIndex, index) {
+			k.Swap(store, index, leftIndex)
+			k.HeapifyDown(ctx, store, leftIndex)
+		}
+	}
+}
+
+// Swap swaps the elements at the given indices.
+func (k Keeper) Swap(
+	store prefix.Store,
+	index1 uint32,
+	index2 uint32,
+) {
+	// No-op case
+	if index1 == index2 {
+		return
+	}
+
+	first := k.MustGetSubaccountAtIndex(store, index1)
+	second := k.MustGetSubaccountAtIndex(store, index2)
+	k.SetSubaccountAtIndex(store, first, index2)
+	k.SetSubaccountAtIndex(store, second, index1)
+}
+
+// Less returns true if the element at the first index is less than
+// the element at the second index.
+func (k Keeper) Less(
+	ctx sdk.Context,
+	store prefix.Store,
+	first uint32,
+	second uint32,
+) bool {
+	firstSubaccountId := k.MustGetSubaccountAtIndex(store, first)
+	secondSubaccountId := k.MustGetSubaccountAtIndex(store, second)
+
+	firstRisk, err := k.GetNetCollateralAndMarginRequirements(
+		ctx,
+		types.Update{
+			SubaccountId: firstSubaccountId,
+		},
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	secondRisk, err := k.GetNetCollateralAndMarginRequirements(
+		ctx,
+		types.Update{
+			SubaccountId: secondSubaccountId,
+		},
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	// Compare the risks of the two subaccounts and sort
+	// them in descending order.
+	return firstRisk.Cmp(secondRisk) > 0
+}

protocol/x/subaccounts/keeper/safety_heap_test.go

@@ -0,0 +1,68 @@
+package keeper_test
+
+import (
+	"math/big"
+	"math/rand"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/types"
+	"github.com/dydxprotocol/v4-chain/protocol/app/config"
+	"github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
+	keepertest "github.com/dydxprotocol/v4-chain/protocol/testutil/keeper"
+	testutil "github.com/dydxprotocol/v4-chain/protocol/testutil/util"
+	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSafetyHeapInsertRemoval(t *testing.T) {
+
+	allSubaccounts := make([]satypes.Subaccount, 0)
+	for i := 0; i < 1000; i++ {
+		subaccount := satypes.Subaccount{
+			Id: &satypes.SubaccountId{
+				Owner: types.MustBech32ifyAddressBytes(
+					config.Bech32PrefixAccAddr,
+					constants.AliceAccAddress,
+				),
+				Number: uint32(i),
+			},
+			AssetPositions: testutil.CreateUsdcAssetPositions(
+				// Create asset positions with balances ranging from -500 to 500.
+				big.NewInt(int64(i - 500)),
+			),
+		}
+
+		// Handle special case.
+		if i == 500 {
+			subaccount.AssetPositions = nil
+		}
+
+		allSubaccounts = append(allSubaccounts, subaccount)
+	}
+
+	for iter := 0; iter < 100; iter++ {
+		// Setup keeper state and test parameters.
+		ctx, subaccountsKeeper, _, _, _, _, _, _, _, _ := keepertest.SubaccountsKeepers(t, false)
+
+		rand.Shuffle(len(allSubaccounts), func(i, j int) {
+			allSubaccounts[i], allSubaccounts[j] = allSubaccounts[j], allSubaccounts[i]
+		})
+
+		store := subaccountsKeeper.GetSafetyHeapStore(ctx, 0, satypes.Long)
+		for _, subaccount := range allSubaccounts {
+			subaccountsKeeper.SetSubaccount(ctx, subaccount)
+			subaccountsKeeper.Insert(ctx, store, *subaccount.Id)
+		}
+
+		// Make sure subaccounts are sorted correctly.
+		for i := 0; i < 1000; i++ {
+			subaccountId := subaccountsKeeper.MustGetSubaccountAtIndex(store, uint32(0))
+
+			// Subaccounts should be sorted by asset position balance.
+			require.Equal(t, uint32(i), subaccountId.Number)
+
+			// Remove the subaccount from the heap.
+			subaccountsKeeper.MustRemoveElementAtIndex(ctx, store, 0)
+		}
+	}
+}