Merge pull request #26 from sei-protocol/tony-chen-size-cmd

Add command that print size of tree

cmd/iaviewer/main.go

@@ -22,8 +22,8 @@ const (
 
 func main() {
 	args := os.Args[1:]
-	if len(args) < 3 || (args[0] != "data" && args[0] != "shape" && args[0] != "versions") {
-		fmt.Fprintln(os.Stderr, "Usage: iaviewer <data|shape|versions> <leveldb dir> <prefix> [version number]")
+	if len(args) < 3 || (args[0] != "data" && args[0] != "shape" && args[0] != "versions" && args[0] != "size") {
+		fmt.Fprintln(os.Stderr, "Usage: iaviewer <data|shape|versions|size> <leveldb dir> <prefix> [version number]")
 		fmt.Fprintln(os.Stderr, "<prefix> is the prefix of db, and the iavl tree of different modules in cosmos-sdk uses ")
 		fmt.Fprintln(os.Stderr, "different <prefix> to identify, just like \"s/k:gov/\" represents the prefix of gov module")
 		os.Exit(1)
@@ -59,6 +59,8 @@ func main() {
 		PrintShape(tree)
 	case "versions":
 		PrintVersions(tree)
+	case "size":
+		PrintSize(tree)
 	}
 }
 
@@ -186,3 +188,24 @@ func PrintVersions(tree *iavl.MutableTree) {
 		fmt.Printf("  %d\n", v)
 	}
 }
+
+func PrintSize(tree *iavl.MutableTree) {
+	count, totalKeySize, totalValueSize := 0, 0, 0
+	keySizeByPrefix, valSizeByPrefix := map[byte]int{}, map[byte]int{}
+	tree.Iterate(func(key []byte, value []byte) bool {
+		count += 1
+		totalKeySize += len(key)
+		totalValueSize += len(value)
+		if _, ok := keySizeByPrefix[key[0]]; !ok {
+			keySizeByPrefix[key[0]] = 0
+			valSizeByPrefix[key[0]] = 0
+		}
+		keySizeByPrefix[key[0]] += len(key)
+		valSizeByPrefix[key[0]] += len(value)
+		return false
+	})
+	fmt.Printf("Total entry count: %d. Total key bytes: %d. Total value bytes: %d\n", count, totalKeySize, totalValueSize)
+	for p := range keySizeByPrefix {
+		fmt.Printf("prefix %d has key bytes %d and value bytes %d\n", p, keySizeByPrefix[p], valSizeByPrefix[p])
+	}
+}