From 8618f216f8bebfbd4b271e4cbb05c1cdedd85d96 Mon Sep 17 00:00:00 2001
From: Geoff Stuart <geoff.vball@gmail.com>
Date: Thu, 13 Oct 2022 15:10:41 -0400
Subject: [PATCH] Fix counting of allocations in nv17 migration test

---
 cmd/lotus-shed/migrations.go | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/cmd/lotus-shed/migrations.go b/cmd/lotus-shed/migrations.go
index 04e63d130..058fbacd5 100644
--- a/cmd/lotus-shed/migrations.go
+++ b/cmd/lotus-shed/migrations.go
@@ -6,6 +6,8 @@ import (
 	"io"
 	"time"
 
+	cbg "github.com/whyrusleeping/cbor-gen"
+
 	"github.com/ipfs/go-cid"
 	"github.com/urfave/cli/v2"
 	"golang.org/x/xerrors"
@@ -339,7 +341,7 @@ func checkPendingVerifiedDeals(stateTreeV8 state.StateTree, stateTreeV9 state.St
 		return xerrors.Errorf("number of allocation IDsf: %d did not match the number of pending verified deals: %d", numAllocationIds, numPendingVerifiedDeals)
 	}
 
-	numAllocations, err := countAllocations(verifregActorV9)
+	numAllocations, err := countAllocations(verifregStateV9, actorStore)
 	if err != nil {
 		return err
 	}
@@ -523,19 +525,36 @@ func checkMinerUnsealedCID(act *types.Actor, stateTreeV9 state.StateTree, store
 	return nil
 }
 
-func countAllocations(verifregState verifreg.State) (int, error) {
-	var count int
-	err := verifregState.ForEachClient(func(addr address.Address, dcap abi.StoragePower) error {
-		allocations, err := verifregState.GetAllocations(addr)
+func countAllocations(verifregState verifreg9.State, store adt.Store) (int, error) {
+	var count = 0
+
+	actorToHamtMap, err := adt9.AsMap(store, verifregState.Allocations, builtin.DefaultHamtBitwidth)
+	if err != nil {
+		return 0, xerrors.Errorf("couldn't get outer map: %x", err)
+	}
+
+	var innerHamtCid cbg.CborCid
+	err = actorToHamtMap.ForEach(&innerHamtCid, func(key string) error {
+		innerMap, err := adt9.AsMap(store, cid.Cid(innerHamtCid), builtin.DefaultHamtBitwidth)
 		if err != nil {
-			return err
+			return xerrors.Errorf("couldn't get outer map: %x", err)
 		}
-		count += len(allocations)
+
+		var allocation verifreg9.Allocation
+		err = innerMap.ForEach(&allocation, func(key string) error {
+			count++
+			return nil
+		})
+		if err != nil {
+			return xerrors.Errorf("couldn't iterate over inner map: %x", err)
+		}
+
 		return nil
 	})
 	if err != nil {
-		return 0, err
+		return 0, xerrors.Errorf("couldn't iterate over outer map: %x", err)
 	}
+
 	return count, nil
 }