From ce65cdd1e5e76cf5070b8db530022e927492f634 Mon Sep 17 00:00:00 2001
From: Peter Rabbitson <ribasushi@protocol.ai>
Date: Wed, 5 Oct 2022 21:10:19 +0200
Subject: [PATCH] Teach shed/sim to understand --tipset=@nnn notation

---
 cli/state.go             | 23 +++++++++++++++++++++++
 cmd/lotus-shed/export.go | 20 +++-----------------
 cmd/lotus-sim/create.go  |  7 +------
 3 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/cli/state.go b/cli/state.go
index c88df8da2..de358815f 100644
--- a/cli/state.go
+++ b/cli/state.go
@@ -44,6 +44,7 @@ import (
 	"github.com/filecoin-project/lotus/chain/consensus/filcns"
 	"github.com/filecoin-project/lotus/chain/state"
 	"github.com/filecoin-project/lotus/chain/stmgr"
+	"github.com/filecoin-project/lotus/chain/store"
 	"github.com/filecoin-project/lotus/chain/types"
 )
 
@@ -294,6 +295,28 @@ func ParseTipSetRef(ctx context.Context, api v0api.FullNode, tss string) (*types
 	return ts, nil
 }
 
+func ParseTipSetRefOffline(ctx context.Context, cs *store.ChainStore, tss string) (*types.TipSet, error) {
+	switch {
+
+	case tss == "" || tss == "@head":
+		return cs.GetHeaviestTipSet(), nil
+
+	case tss[0] != '@':
+		cids, err := ParseTipSetString(tss)
+		if err != nil {
+			return nil, xerrors.Errorf("failed to parse tipset (%q): %w", tss, err)
+		}
+		return cs.LoadTipSet(ctx, types.NewTipSetKey(cids...))
+
+	default:
+		var h uint64
+		if _, err := fmt.Sscanf(tss, "@%d", &h); err != nil {
+			return nil, xerrors.Errorf("parsing height tipset ref: %w", err)
+		}
+		return cs.GetTipsetByHeight(ctx, abi.ChainEpoch(h), cs.GetHeaviestTipSet(), true)
+	}
+}
+
 var StatePowerCmd = &cli.Command{
 	Name:      "power",
 	Usage:     "Query network or miner power",
diff --git a/cmd/lotus-shed/export.go b/cmd/lotus-shed/export.go
index 8de3f48df..459de3383 100644
--- a/cmd/lotus-shed/export.go
+++ b/cmd/lotus-shed/export.go
@@ -32,7 +32,6 @@ import (
 
 	"github.com/filecoin-project/lotus/blockstore"
 	"github.com/filecoin-project/lotus/chain/store"
-	"github.com/filecoin-project/lotus/chain/types"
 	lcli "github.com/filecoin-project/lotus/cli"
 	"github.com/filecoin-project/lotus/cmd/lotus-shed/shedgen"
 	"github.com/filecoin-project/lotus/node/repo"
@@ -125,22 +124,9 @@ var exportChainCmd = &cli.Command{
 		fullstate := cctx.Bool("full-state")
 		skipoldmsgs := cctx.Bool("skip-old-msgs")
 
-		var ts *types.TipSet
-		if tss := cctx.String("tipset"); tss != "" {
-			cids, err := lcli.ParseTipSetString(tss)
-			if err != nil {
-				return xerrors.Errorf("failed to parse tipset (%q): %w", tss, err)
-			}
-
-			tsk := types.NewTipSetKey(cids...)
-
-			selts, err := cs.LoadTipSet(context.Background(), tsk)
-			if err != nil {
-				return xerrors.Errorf("loading tipset: %w", err)
-			}
-			ts = selts
-		} else {
-			ts = cs.GetHeaviestTipSet()
+		ts, err := lcli.ParseTipSetRefOffline(ctx, cs, cctx.String("tipset"))
+		if err != nil {
+			return err
 		}
 
 		if fullstate {
diff --git a/cmd/lotus-sim/create.go b/cmd/lotus-sim/create.go
index 23ea454a3..bc7f9ade9 100644
--- a/cmd/lotus-sim/create.go
+++ b/cmd/lotus-sim/create.go
@@ -31,12 +31,7 @@ var createSimCommand = &cli.Command{
 			}
 			ts = node.Chainstore.GetHeaviestTipSet()
 		case 1:
-			cids, err := lcli.ParseTipSetString(cctx.Args().Get(1))
-			if err != nil {
-				return err
-			}
-			tsk := types.NewTipSetKey(cids...)
-			ts, err = node.Chainstore.LoadTipSet(cctx.Context, tsk)
+			ts, err = lcli.ParseTipSetRefOffline(cctx.Context, node.Chainstore, cctx.Args().Get(1))
 			if err != nil {
 				return err
 			}