Update snapshot size implementation

Apply further code changes provided by @dougm.

- credit: @dougm
- refs vmware#2243

govc/vm/snapshot/size.go

@@ -39,8 +39,7 @@ func removeKey(l *[]int, key int) {
 */
 func SnapshotSize(
 
-	// info == vm.Snapshot
-	info *types.VirtualMachineSnapshotInfo,
+	info types.ManagedObjectReference,
 
 	// this function cannot be expected to process a snapshot tree. Instead,
 	// something else looks at the tree and calls this for each snapshot. This
@@ -73,7 +72,7 @@ func SnapshotSize(
 
 		// if the layout MOID matches active snapshot MOID
 		// Q: Why match on this? Is this because the snapshot is growing?
-		if layout.Key.Value == info.CurrentSnapshot.Value {
+		if layout.Key.Value == info.Value {
 
 			// gather file keys if snapshot is current?
 			fileKeyList = append(fileKeyList, int(layout.DataKey)) // The .vmsn file

govc/vm/snapshot/tree.go

@@ -39,8 +39,10 @@ type tree struct {
 	description bool
 	fullPath    bool
 	id          bool
+	size        bool
 
-	info *types.VirtualMachineSnapshotInfo
+	info   *types.VirtualMachineSnapshotInfo
+	layout *types.VirtualMachineFileLayoutEx
 }
 
 func init() {
@@ -60,6 +62,7 @@ func (cmd *tree) Register(ctx context.Context, f *flag.FlagSet) {
 	f.BoolVar(&cmd.fullPath, "f", false,
 		"Print the full path prefix for snapshot")
 	f.BoolVar(&cmd.id, "i", false, "Print the snapshot id")
+	f.BoolVar(&cmd.size, "s", false, "Print the snapshot size")
 }
 
 func (cmd *tree) Description() string {
@@ -79,7 +82,7 @@ func (cmd *tree) Process(ctx context.Context) error {
 	return nil
 }
 
-func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapshotTree) {
+func (cmd *tree) write(level int, parent string, pref *types.ManagedObjectReference, st []types.VirtualMachineSnapshotTree) {
 	for _, s := range st {
 		sname := s.Name
 
@@ -93,7 +96,10 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			names = append(names, sname)
 		}
 
+		isCurrent := false
+
 		if s.Snapshot == *cmd.info.CurrentSnapshot {
+			isCurrent = true
 			if cmd.current {
 				names = append(names, ".")
 			} else if cmd.currentName {
@@ -106,6 +112,12 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			var attr []string
 			var meta string
 
+			if cmd.size {
+				size := SnapshotSize(s.Snapshot, pref, cmd.layout, isCurrent)
+
+				attr = append(attr, units.ByteSize(size).String())
+			}
+
 			if cmd.id {
 				attr = append(attr, s.Snapshot.Value)
 			}
@@ -128,7 +140,7 @@ func (cmd *tree) write(level int, parent string, st []types.VirtualMachineSnapsh
 			}
 		}
 
-		cmd.write(level+2, sname, s.ChildSnapshotList)
+		cmd.write(level+2, sname, &s.Snapshot, s.ChildSnapshotList)
 	}
 }
 
@@ -162,10 +174,11 @@ func (cmd *tree) Run(ctx context.Context, f *flag.FlagSet) error {
 	}
 
 	cmd.info = o.Snapshot
-	size := SnapshotSize(o.Snapshot, nil, o.LayoutEx, true)
+	cmd.layout = o.LayoutEx
+	size := SnapshotSize(o.Reference(), nil, o.LayoutEx, true)
 	fmt.Printf("size=%s\n", units.ByteSize(size))
 
-	cmd.write(0, "", o.Snapshot.RootSnapshotList)
+	cmd.write(0, "", nil, o.Snapshot.RootSnapshotList)
 
 	return nil
 }