fix(protocol-designer): fix slots availability functions to support multiple modules

protocol-designer/src/pages/CreateNewProtocolWizard/SelectModules.tsx

@@ -35,7 +35,13 @@ import {
   FLEX_SUPPORTED_MODULE_MODELS,
   OT2_SUPPORTED_MODULE_MODELS,
 } from './constants'
-import { getNumOptions, getNumSlotsAvailable } from './utils'
+import {
+  getNumOptionsForModules,
+  getNumSlotsAvailable,
+  getModuleDistribution,
+  getAvailableSlots,
+  getCanAddModule,
+} from './utils'
 import { HandleEnter } from './HandleEnter'
 
 import type { DropdownBorder } from '@opentrons/components'
@@ -44,7 +50,6 @@ import type { FormModule, FormModules } from '../../step-forms'
 import type { WizardTileProps } from './types'
 
 const MAX_MAGNETIC_BLOCKS = 4
-const MAGNETIC_BLOCKS_ADJUSTMENT = 3
 
 export function SelectModules(props: WizardTileProps): JSX.Element | null {
   const { goBack, proceed, watch, setValue } = props
@@ -60,6 +65,11 @@ export function SelectModules(props: WizardTileProps): JSX.Element | null {
       ? FLEX_SUPPORTED_MODULE_MODELS
       : OT2_SUPPORTED_MODULE_MODELS
 
+  const distribution = getModuleDistribution(modules)
+  console.log('distribution', distribution)
+  const availableSlots = getAvailableSlots(distribution)
+  console.log('availableSlots', availableSlots)
+
   const numSlotsAvailable = getNumSlotsAvailable(modules, additionalEquipment)
   const hasNoAvailableSlots = numSlotsAvailable === 0
   const numMagneticBlocks =
@@ -86,20 +96,23 @@ export function SelectModules(props: WizardTileProps): JSX.Element | null {
   ]
 
   const handleAddModule = (moduleModel: ModuleModel): void => {
-    if (hasNoAvailableSlots) {
-      makeSnackbar(t('slots_limit_reached') as string)
-    } else {
+    const moduleType = getModuleType(moduleModel)
+    const distribution = getModuleDistribution(modules)
+
+    if (getCanAddModule(moduleType, distribution)) {
       setValue('modules', {
         ...modules,
         [uuid()]: {
           model: moduleModel,
-          type: getModuleType(moduleModel),
+          type: moduleType,
           slot:
             robotType === FLEX_ROBOT_TYPE
               ? DEFAULT_SLOT_MAP_FLEX[moduleModel]
               : DEFAULT_SLOT_MAP_OT2[getModuleType(moduleModel)],
         },
       })
+    } else {
+      makeSnackbar(t('cannot_add_module') as string)
     }
   }
 
@@ -254,12 +267,9 @@ export function SelectModules(props: WizardTileProps): JSX.Element | null {
                           )
                         },
                         dropdownType: 'neutral' as DropdownBorder,
-                        filterOptions: getNumOptions(
-                          module.model === 'magneticBlockV1'
-                            ? numSlotsAvailable +
-                                MAGNETIC_BLOCKS_ADJUSTMENT +
-                                module.count
-                            : numSlotsAvailable + module.count
+                        filterOptions: getNumOptionsForModules(
+                          module.type,
+                          distribution
                         ),
                       }
                       return (

protocol-designer/src/pages/CreateNewProtocolWizard/utils.tsx

@@ -1,10 +1,12 @@
 import {
-  MAGNETIC_BLOCK_TYPE,
-  STAGING_AREA_CUTOUTS,
-  THERMOCYCLER_MODULE_TYPE,
   getLabwareDefURI,
   getLabwareDisplayName,
   getPipetteSpecsV2,
+  HEATERSHAKER_MODULE_TYPE,
+  MAGNETIC_BLOCK_TYPE,
+  STAGING_AREA_CUTOUTS,
+  TEMPERATURE_MODULE_TYPE,
+  THERMOCYCLER_MODULE_TYPE,
   WASTE_CHUTE_CUTOUT,
 } from '@opentrons/shared-data'
 import wasteChuteImage from '../../assets/images/waste_chute.png'
@@ -15,9 +17,11 @@ import type {
   LabwareDefByDefURI,
   LabwareDefinition2,
   PipetteName,
+  ModuleType,
 } from '@opentrons/shared-data'
 import type { DropdownOption } from '@opentrons/components'
 import type { AdditionalEquipment, WizardFormState } from './types'
+import type { FormModules } from '../../step-forms'
 
 const TOTAL_MODULE_SLOTS = 8
 const MIDDLE_SLOT_NUM = 4
@@ -29,6 +33,48 @@ export const getNumOptions = (length: number): DropdownOption[] => {
   }))
 }
 
+export const getNumOptionsForModules = (
+  moduleType: ModuleType,
+  distribution: ModuleDistribution
+): DropdownOption[] => {
+  const { tc, hs, mb, tm } = distribution
+  let maxCount: number
+
+  const totalOccupiedSlots = tc * 2 + hs + tm + mb
+  const maxSlots = mb === 0 ? 7 : 11
+  const availableSlots = maxSlots - totalOccupiedSlots
+
+  switch (moduleType) {
+    case THERMOCYCLER_MODULE_TYPE:
+      maxCount = tc === 0 ? 1 : 0
+      break
+    case MAGNETIC_BLOCK_TYPE:
+      maxCount = Math.min(11 - (tc * 2 + hs + tm), 11)
+      break
+    case HEATERSHAKER_MODULE_TYPE:
+    case TEMPERATURE_MODULE_TYPE:
+      if (tc === 0) {
+        maxCount = Math.min(
+          7,
+          availableSlots + (moduleType === HEATERSHAKER_MODULE_TYPE ? hs : tm)
+        )
+      } else {
+        maxCount = Math.min(
+          5,
+          availableSlots + (moduleType === HEATERSHAKER_MODULE_TYPE ? hs : tm)
+        )
+      }
+      break
+    default:
+      maxCount = 0
+  }
+
+  return Array.from({ length: maxCount }, (_, i) => ({
+    name: `${i + 1}`,
+    value: `${i + 1}`,
+  }))
+}
+
 export const getNumSlotsAvailable = (
   modules: WizardFormState['modules'],
   additionalEquipment: WizardFormState['additionalEquipment']
@@ -49,7 +95,7 @@ export const getNumSlotsAvailable = (
     filteredModuleLength = filteredModuleLength + 1
   }
   if (magneticBlocks.length > 0) {
-    //  once blocks exceed 4, then we dont' want to subtract the amount available
+    //  once blocks exceed 4, then we don't want to subtract the amount available
     //  because block can go into the center slots where all other modules/trashes can not
     const numBlocks =
       magneticBlocks.length > 4 ? MIDDLE_SLOT_NUM : magneticBlocks.length
@@ -237,3 +283,83 @@ export const getTrashSlot = (values: WizardFormState): string => {
   }
   return unoccupiedSlot?.value
 }
+
+interface ModuleDistribution {
+  tc: number
+  hs: number
+  mb: number
+  tm: number
+}
+
+const TOTAL_SLOTS_WITHOUT_TWO_COL = 7
+const TWO_COL_SLOTS = 4
+
+export const getModuleDistribution = (
+  modules: FormModules | null
+): ModuleDistribution => {
+  let tc = 0
+  let hs = 0
+  let mb = 0
+  let tm = 0
+
+  if (modules === null) return { tc, hs, mb, tm }
+
+  Object.values(modules).forEach(module => {
+    switch (module.type) {
+      case THERMOCYCLER_MODULE_TYPE:
+        // TC occupies A1+B1
+        tc++
+        break
+      case HEATERSHAKER_MODULE_TYPE:
+        hs++
+        break
+      case MAGNETIC_BLOCK_TYPE:
+        mb++
+        break
+      case TEMPERATURE_MODULE_TYPE:
+        tm++
+        break
+    }
+  })
+
+  return { tc, hs, mb, tm }
+}
+
+export const getAvailableSlots = (
+  distribution: ModuleDistribution
+): {
+  regular: number
+  magnetic: number
+} => {
+  if (distribution === null)
+    return { regular: TOTAL_SLOTS_WITHOUT_TWO_COL, magnetic: TWO_COL_SLOTS }
+  const { tc, hs, mb, tm } = distribution
+  const availableMagneticSlots = 11 - (tc * 2 + hs + tm + mb)
+  const availableRegularSlots = 7 - (hs + tm + (tc > 0 ? 2 : 0))
+
+  return {
+    regular: Math.max(0, availableRegularSlots),
+    magnetic: Math.max(0, availableMagneticSlots),
+  }
+}
+
+export const getCanAddModule = (
+  moduleType: ModuleType,
+  distribution: ModuleDistribution
+): boolean => {
+  const availableSlots = getAvailableSlots(distribution)
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  const { tc, hs, mb, tm } = distribution
+
+  switch (moduleType) {
+    case THERMOCYCLER_MODULE_TYPE:
+      return tc === 0 && availableSlots.regular >= 2
+    case MAGNETIC_BLOCK_TYPE:
+      return availableSlots.magnetic > 0
+    case HEATERSHAKER_MODULE_TYPE:
+    case TEMPERATURE_MODULE_TYPE:
+      return availableSlots.regular > 0 || (hs + tm < 5 && tc === 0)
+    default:
+      return false
+  }
+}