tried to get it to work

src/css/ingame_hud/buildings_toolbar.scss

@@ -109,6 +109,14 @@
                     }
                 }
             }
+
+            &.puzzle-lock {
+                display: inline;
+                background-image: uiResource("res/ui/unlocked-building.png");
+                background-size: 50%;
+                background-position-x: right;
+                background-position-y: top;
+            }
         }
     }
 }

src/js/game/hud/parts/base_toolbar.js

@@ -35,6 +35,7 @@ export class HUDBaseToolbar extends BaseHUDPart {
          * selected: boolean,
          * element: HTMLElement,
          * index: number
+         * puzzleLocked: boolean,
          * }>} */
         this.buildingHandles = {};
     }
@@ -105,19 +106,27 @@ export class HUDBaseToolbar extends BaseHUDPart {
             );
             itemContainer.setAttribute("data-icon", "building_icons/" + metaBuilding.getId() + ".png");
             itemContainer.setAttribute("data-id", metaBuilding.getId());
-
             binding.add(() => this.selectBuildingForPlacement(metaBuilding));
 
             this.trackClicks(itemContainer, () => this.selectBuildingForPlacement(metaBuilding), {
                 clickSound: null,
             });
 
+            //new stuff
+            const puzzleLock = makeDiv(itemContainer, null, ["puzzle-lock"]);
+            puzzleLock.setAttribute("data-icon", "unlocked_building.png");
+
+            this.trackClicks(puzzleLock, () => this.togglePuzzleLock(metaBuilding), {
+                clickSound: null,
+            });
+
             this.buildingHandles[metaBuilding.id] = {
                 metaBuilding,
                 element: itemContainer,
                 unlocked: false,
                 selected: false,
                 index: i,
+                puzzleLocked: false,
             };
         }
 
@@ -247,4 +256,24 @@ export class HUDBaseToolbar extends BaseHUDPart {
         this.root.hud.signals.buildingSelectedForPlacement.dispatch(metaBuilding);
         this.onSelectedPlacementBuildingChanged(metaBuilding);
     }
+
+    /**
+     * @param {MetaBuilding} metaBuilding
+     */
+    togglePuzzleLock(metaBuilding) {
+        if (!this.visibilityCondition()) {
+            // Not active
+            return;
+        }
+
+        if (!metaBuilding.getIsUnlocked(this.root)) {
+            this.root.soundProxy.playUiError();
+            return STOP_PROPAGATION;
+        }
+
+        const handle = this.buildingHandles[metaBuilding.getId()];
+        handle.puzzleLocked = !handle.puzzleLocked;
+
+        this.root.soundProxy.playUiClick();
+    }
 }

src/js/game/modes/puzzle.js

@@ -7,6 +7,9 @@ import { types } from "../../savegame/serialization";
 import { enumGameModeTypes, GameMode } from "../game_mode";
 import { HUDPuzzleBackToMenu } from "../hud/parts/puzzle_back_to_menu";
 import { HUDPuzzleDLCLogo } from "../hud/parts/puzzle_dlc_logo";
+import { gMetaBuildingRegistry } from "../../core/global_registries";
+import { MetaBalancerBuilding } from "../buildings/balancer";
+import { MetaUndergroundBeltBuilding } from "../buildings/underground_belt";
 
 export class PuzzleGameMode extends GameMode {
     static getType() {
@@ -36,6 +39,13 @@ export class PuzzleGameMode extends GameMode {
         this.zoneHeight = data.zoneHeight || 6;
     }
 
+    /**
+     * @param {typeof import("../meta_building").MetaBuilding} building
+     */
+    isBuildingExcluded(building) {
+        return this.hiddenBuildings.indexOf(building) >= 0;
+    }
+
     getSaveData() {
         const save = this.root.savegame.getCurrentDump();
         if (!save) {

src/js/game/modes/puzzle_play.js

@@ -28,6 +28,7 @@ import { createLogger } from "../../core/logging";
 import { HUDPuzzleCompleteNotification } from "../hud/parts/puzzle_complete_notification";
 import { HUDPuzzlePlaySettings } from "../hud/parts/puzzle_play_settings";
 import { MetaBlockBuilding } from "../buildings/block";
+import { gMetaBuildingRegistry } from "../../core/global_registries";
 import { MetaBuilding } from "../meta_building";
 
 const logger = createLogger("puzzle-play");
@@ -46,8 +47,8 @@ export class PuzzlePlayGameMode extends PuzzleGameMode {
     constructor(root, { puzzle }) {
         super(root);
 
-        /** @type {(typeof MetaBuilding)[]} */
+        /** @type {Array<typeof MetaBuilding>} */
-        this.hiddenBuildings = [
+        const excludedBuildings = [
             MetaConstantProducerBuilding,
             MetaGoalAcceptorBuilding,
             MetaBlockBuilding,
@@ -69,7 +70,8 @@ export class PuzzlePlayGameMode extends PuzzleGameMode {
             MetaComparatorBuilding,
             MetaTransistorBuilding,
         ];
-        this.hiddenBuildings.push(puzzle.game.excludedBuildings);
+
+        this.hiddenBuildings = excludedBuildings.concat(puzzle.game.excludedBuildings);
 
         this.additionalHudParts.puzzlePlayMetadata = HUDPuzzlePlayMetadata;
         this.additionalHudParts.puzzlePlaySettings = HUDPuzzlePlaySettings;

src/js/game/modes/regular.js

@@ -1,5 +1,6 @@
 /* typehints:start */
 import { GameRoot } from "../root";
+import { MetaBuilding } from "../meta_building";
 /* typehints:end */
 
 import { findNiceIntegerValue } from "../../core/utils";
@@ -582,6 +583,7 @@ export class RegularGameMode extends GameMode {
             this.additionalHudParts.sandboxController = HUDSandboxController;
         }
 
+        /** @type {(typeof MetaBuilding)[]} */
         this.hiddenBuildings = [MetaConstantProducerBuilding, MetaGoalAcceptorBuilding, MetaBlockBuilding];
     }
 

src/js/savegame/puzzle_serializer.js

@@ -33,7 +33,6 @@ export class PuzzleSerializer {
          * @type {import("./savegame_typedefs").PuzzleGameData["buildings"]}
          */
         let buildings = [];
-
         for (const entity of root.entityMgr.getAllWithComponent(StaticMapEntityComponent)) {
             const staticComp = entity.components.StaticMapEntity;
             const signalComp = entity.components.ConstantSignal;
@@ -90,6 +89,7 @@ export class PuzzleSerializer {
                 w: mode.zoneWidth,
                 h: mode.zoneHeight,
             },
+            //read from the toolbar when making a puzzle
             excludedBuildings: [],
         };
     }

src/js/savegame/savegame_typedefs.js

@@ -87,7 +87,7 @@ import { MetaBuilding } from "../game/meta_building";
  *   version: number;
  *   bounds: { w: number; h: number; },
  *   buildings: (PuzzleGameBuildingGoal | PuzzleGameBuildingConstantProducer | PuzzleGameBuildingBlock)[],
- *   excludedBuildings: (typeof MetaBuilding)[],
+ *   excludedBuildings: Array<typeof MetaBuilding>,
  * }} PuzzleGameData
  */