Enhance delete foliage layer command to remove masks from affected terrains and improve undo state restoration

src/commands/delete-foliage-layer-command.ts

@@ -1,4 +1,5 @@
 import { createOpaqueId } from "../core/ids";
+import { cloneTerrain, type Terrain } from "../document/terrains";
 import { cloneFoliageLayer, type FoliageLayer } from "../foliage/foliage";
 
 import type { EditorCommand } from "./command";
@@ -7,6 +8,7 @@ export function createDeleteFoliageLayerCommand(
   foliageLayerId: string
 ): EditorCommand {
   let deletedLayer: FoliageLayer | null = null;
+  let previousTerrains: Record<string, Terrain> | null = null;
 
   return {
     id: createOpaqueId("command"),
@@ -23,15 +25,44 @@ export function createDeleteFoliageLayerCommand(
         deletedLayer = cloneFoliageLayer(currentLayer);
       }
 
+      if (previousTerrains === null) {
+        previousTerrains = Object.fromEntries(
+          Object.entries(currentDocument.terrains)
+            .filter(
+              ([, terrain]) => terrain.foliageMasks[foliageLayerId] !== undefined
+            )
+            .map(([terrainId, terrain]) => [terrainId, cloneTerrain(terrain)])
+        );
+      }
+
       const nextFoliageLayers = {
         ...currentDocument.foliageLayers
       };
+      const nextTerrains = {
+        ...currentDocument.terrains
+      };
 
       delete nextFoliageLayers[foliageLayerId];
 
+      for (const [terrainId, terrain] of Object.entries(nextTerrains)) {
+        if (terrain.foliageMasks[foliageLayerId] === undefined) {
+          continue;
+        }
+
+        const nextFoliageMasks = {
+          ...terrain.foliageMasks
+        };
+        delete nextFoliageMasks[foliageLayerId];
+        nextTerrains[terrainId] = cloneTerrain({
+          ...terrain,
+          foliageMasks: nextFoliageMasks
+        });
+      }
+
       context.setDocument({
         ...currentDocument,
-        foliageLayers: nextFoliageLayers
+        foliageLayers: nextFoliageLayers,
+        terrains: nextTerrains
       });
     },
     undo(context) {
@@ -40,13 +71,23 @@ export function createDeleteFoliageLayerCommand(
       }
 
       const currentDocument = context.getDocument();
+      const restoredTerrains = {
+        ...currentDocument.terrains
+      };
+
+      if (previousTerrains !== null) {
+        for (const [terrainId, terrain] of Object.entries(previousTerrains)) {
+          restoredTerrains[terrainId] = cloneTerrain(terrain);
+        }
+      }
 
       context.setDocument({
         ...currentDocument,
         foliageLayers: {
           ...currentDocument.foliageLayers,
           [deletedLayer.id]: cloneFoliageLayer(deletedLayer)
-        }
+        },
+        terrains: restoredTerrains
       });
     }
   };