check dir upgrade even in RemoveChild case

Author: schomatis

io/directory.go

@@ -214,7 +214,9 @@ func (d *BasicDirectory) needsToSwitchToHAMTDir(name string, nodeToAdd ipld.Node
 		}
 		operationSizeChange -= estimatedLinkSize(name, entryToRemove.Cid)
 	}
-	operationSizeChange += estimatedLinkSize(name, nodeToAdd.Cid())
+	if nodeToAdd != nil {
+		operationSizeChange += estimatedLinkSize(name, nodeToAdd.Cid())
+	}
 
 	return d.estimatedSize + operationSizeChange >= HAMTShardingSize, nil
 }
@@ -613,15 +615,38 @@ func (d *UpgradeableDirectory) RemoveChild(ctx context.Context, name string) err
 		if err != nil {
 			return err
 		}
-
-		if switchToBasic {
-			basicDir, err := hamtDir.switchToBasic(ctx)
-			if err != nil {
-				return err
-			}
-			d.Directory = basicDir
+		if !switchToBasic {
+			return hamtDir.RemoveChild(ctx, name)
+		}
+		basicDir, err := hamtDir.switchToBasic(ctx)
+		if err != nil {
+			return err
 		}
+		err = basicDir.RemoveChild(ctx, name)
+		if err != nil {
+			return err
+		}
+		d.Directory = basicDir
+		return nil
 	}
 
-	return d.Directory.RemoveChild(ctx, name)
+	// BasicDirectory
+	basicDir := d.Directory.(*BasicDirectory)
+	switchToHAMT, err := basicDir.needsToSwitchToHAMTDir(name, nil)
+	if err != nil {
+		return err
+	}
+	if !switchToHAMT {
+		return basicDir.RemoveChild(ctx, name)
+	}
+	hamtDir, err = basicDir.SwitchToSharding(ctx)
+	if err != nil {
+		return err
+	}
+	hamtDir.RemoveChild(ctx, name)
+	if err != nil {
+		return err
+	}
+	d.Directory = hamtDir
+	return nil
 }