From d9131f756393b85cda42e559a1650d882cf114c9 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 7 Aug 2024 08:23:51 +1000
Subject: [PATCH] feat(ui): restore object count to layer titles

---
 .../components/common/CanvasEntityTitle.tsx   | 22 ++----------
 .../hooks/useEntityObjectCount.ts             | 28 +++++++++++++++
 .../controlLayers/hooks/useEntityTitle.ts     | 36 +++++++++++++++++++
 3 files changed, 67 insertions(+), 19 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/controlLayers/hooks/useEntityObjectCount.ts
 create mode 100644 invokeai/frontend/web/src/features/controlLayers/hooks/useEntityTitle.ts

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityTitle.tsx b/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityTitle.tsx
index 3853f7249c..79e0677cdc 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityTitle.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/common/CanvasEntityTitle.tsx
@@ -1,29 +1,13 @@
 import { Text } from '@invoke-ai/ui-library';
 import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
 import { useEntityIsSelected } from 'features/controlLayers/hooks/useEntityIsSelected';
-import { memo, useMemo } from 'react';
-import { useTranslation } from 'react-i18next';
-import { assert } from 'tsafe';
+import { useEntityTitle } from 'features/controlLayers/hooks/useEntityTitle';
+import { memo } from 'react';
 
 export const CanvasEntityTitle = memo(() => {
-  const { t } = useTranslation();
   const entityIdentifier = useEntityIdentifierContext();
   const isSelected = useEntityIsSelected(entityIdentifier);
-  const title = useMemo(() => {
-    if (entityIdentifier.type === 'inpaint_mask') {
-      return t('controlLayers.inpaintMask');
-    } else if (entityIdentifier.type === 'control_adapter') {
-      return t('controlLayers.globalControlAdapter');
-    } else if (entityIdentifier.type === 'layer') {
-      return t('controlLayers.layer');
-    } else if (entityIdentifier.type === 'ip_adapter') {
-      return t('controlLayers.ipAdapter');
-    } else if (entityIdentifier.type === 'regional_guidance') {
-      return t('controlLayers.regionalGuidance');
-    } else {
-      assert(false, 'Unexpected entity type');
-    }
-  }, [entityIdentifier.type, t]);
+  const title = useEntityTitle(entityIdentifier);
 
   return (
     <Text size="sm" fontWeight="semibold" userSelect="none" color={isSelected ? 'base.100' : 'base.300'}>
diff --git a/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityObjectCount.ts b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityObjectCount.ts
new file mode 100644
index 0000000000..4a9f383b4c
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityObjectCount.ts
@@ -0,0 +1,28 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { useAppSelector } from 'app/store/storeHooks';
+import { selectCanvasV2Slice, selectEntity } from 'features/controlLayers/store/canvasV2Slice';
+import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
+import { useMemo } from 'react';
+
+export const useEntityObjectCount = (entityIdentifier: CanvasEntityIdentifier) => {
+  const selectObjectCount = useMemo(
+    () =>
+      createSelector(selectCanvasV2Slice, (canvasV2) => {
+        const entity = selectEntity(canvasV2, entityIdentifier);
+        if (!entity) {
+          return 0;
+        } else if (entity.type === 'layer') {
+          return entity.objects.length;
+        } else if (entity.type === 'inpaint_mask') {
+          return entity.objects.length;
+        } else if (entity.type === 'regional_guidance') {
+          return entity.objects.length;
+        } else {
+          return 0;
+        }
+      }),
+    [entityIdentifier]
+  );
+  const objectCount = useAppSelector(selectObjectCount);
+  return objectCount;
+};
diff --git a/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityTitle.ts b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityTitle.ts
new file mode 100644
index 0000000000..82b9a457d6
--- /dev/null
+++ b/invokeai/frontend/web/src/features/controlLayers/hooks/useEntityTitle.ts
@@ -0,0 +1,36 @@
+import { useEntityObjectCount } from 'features/controlLayers/hooks/useEntityObjectCount';
+import type { CanvasEntityIdentifier } from 'features/controlLayers/store/types';
+import { useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { assert } from 'tsafe';
+
+export const useEntityTitle = (entityIdentifier: CanvasEntityIdentifier) => {
+  const { t } = useTranslation();
+
+  const objectCount = useEntityObjectCount(entityIdentifier);
+
+  const title = useMemo(() => {
+    const parts: string[] = [];
+    if (entityIdentifier.type === 'inpaint_mask') {
+      parts.push(t('controlLayers.inpaintMask'));
+    } else if (entityIdentifier.type === 'control_adapter') {
+      parts.push(t('controlLayers.globalControlAdapter'));
+    } else if (entityIdentifier.type === 'layer') {
+      parts.push(t('controlLayers.layer'));
+    } else if (entityIdentifier.type === 'ip_adapter') {
+      parts.push(t('controlLayers.ipAdapter'));
+    } else if (entityIdentifier.type === 'regional_guidance') {
+      parts.push(t('controlLayers.regionalGuidance'));
+    } else {
+      assert(false, 'Unexpected entity type');
+    }
+
+    if (objectCount > 0) {
+      parts.push(`(${objectCount})`);
+    }
+
+    return parts.join(' ');
+  }, [entityIdentifier.type, objectCount, t]);
+
+  return title;
+};