From ea53793162e957ee3cd5ac59bf490680a577e1a5 Mon Sep 17 00:00:00 2001
From: Hongtao Lye <hongtao.lye@toeverything.info>
Date: Mon, 9 Dec 2024 22:58:54 +0800
Subject: [PATCH] feat: expand button

---
 .../src/renderer/elements/mindmap.ts          |  10 +-
 .../model/src/elements/mindmap/mindmap.ts     | 121 ++++++++----------
 .../model/src/elements/mindmap/style.ts       |   6 +-
 .../affine/model/src/elements/shape/shape.ts  |   2 +-
 4 files changed, 64 insertions(+), 75 deletions(-)

diff --git a/packages/affine/block-surface/src/renderer/elements/mindmap.ts b/packages/affine/block-surface/src/renderer/elements/mindmap.ts
index 2da67ebb27034..6058c46a64704 100644
--- a/packages/affine/block-surface/src/renderer/elements/mindmap.ts
+++ b/packages/affine/block-surface/src/renderer/elements/mindmap.ts
@@ -23,13 +23,17 @@ export function mindmap(
 
   model.traverse((to, from) => {
     if (from) {
-      const connector = model.getConnector(from, to);
-      if (!connector) return;
+      const result = model.getConnector(from, to);
+      if (!result) return;
 
+      const { connector, outdated } = result;
       const elementGetter = (id: string) =>
         model.surface.getElementById(id) ??
         (model.surface.doc.getBlockById(id) as GfxModel);
-      ConnectorPathGenerator.updatePath(connector, null, elementGetter);
+
+      if (outdated) {
+        ConnectorPathGenerator.updatePath(connector, null, elementGetter);
+      }
 
       const dx = connector.x - bound.x;
       const dy = connector.y - bound.y;
diff --git a/packages/affine/model/src/elements/mindmap/mindmap.ts b/packages/affine/model/src/elements/mindmap/mindmap.ts
index 7c8c21bab0804..a727036c1af51 100644
--- a/packages/affine/model/src/elements/mindmap/mindmap.ts
+++ b/packages/affine/model/src/elements/mindmap/mindmap.ts
@@ -25,7 +25,7 @@ import { DocCollection, type Y } from '@blocksuite/store';
 import { generateKeyBetween } from 'fractional-indexing';
 import { z } from 'zod';
 
-import type { ConnectorStyle, MindmapStyleGetter } from './style.js';
+import type { MindmapStyleGetter } from './style.js';
 
 import { LayoutType, MindmapStyle } from '../../consts/mindmap.js';
 import { LocalConnectorElementModel } from '../connector/local-connector.js';
@@ -251,12 +251,12 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
     const cacheKey = `${from.element.xywh}-${to.element.xywh}-${layout}-${this.style}`;
 
     if (connector.cache.get('MINDMAP_CONNECTOR') === cacheKey) {
-      return { outdated: false, cacheKey };
+      return false;
     } else if (updateKey) {
       connector.cache.set('MINDMAP_CONNECTOR', cacheKey);
     }
 
-    return { outdated: true, cacheKey };
+    return true;
   }
 
   protected override _getXYWH(): Bound {
@@ -271,58 +271,6 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
     noop();
   }
 
-  protected addConnector(
-    from: MindmapNode,
-    to: MindmapNode,
-    layout: LayoutType,
-    connectorStyle: ConnectorStyle
-  ) {
-    const id = `#${from.id}-${to.id}`;
-
-    if (this.connectors.has(id)) {
-      const connector = this.connectors.get(id)!;
-      const { outdated } = this._isConnectorOutdated({
-        connector,
-        from,
-        to,
-        layout,
-      });
-
-      if (!outdated) {
-        return connector;
-      }
-    } else {
-      const connector = new LocalConnectorElementModel(this.surface);
-      // update cache key
-      this._isConnectorOutdated({
-        connector,
-        from,
-        to,
-        layout,
-      });
-      this.connectors.set(id, connector);
-    }
-
-    const connector = this.connectors.get(id)!;
-
-    connector.id = id;
-    connector.source = {
-      id: from.id,
-      position: layout === LayoutType.RIGHT ? [1, 0.5] : [0, 0.5],
-    };
-    connector.target = {
-      id: to.id,
-      position: layout === LayoutType.RIGHT ? [0, 0.5] : [1, 0.5],
-    };
-
-    Object.entries(connectorStyle).forEach(([key, value]) => {
-      // @ts-ignore
-      connector[key as unknown] = value;
-    });
-
-    return connector;
-  }
-
   addNode(
     /**
      * The parent node id of the new node. If it's null, the node will be the root node
@@ -553,7 +501,6 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
       this.collapseButtons.get(id) || new LocalShapeElementModel(this.surface);
 
     if (
-      !btnExisted ||
       this._isCollapseButtonOutdated({
         button: collapseButton,
         node,
@@ -561,9 +508,8 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
       })
     ) {
       const style = this.styleGetter.getNodeStyle(node, this.getPath(node));
-      const buttonStyle = node.detail.collapsed
-        ? style.expandButton
-        : style.collapseButton;
+      const collapsed = node.detail.collapsed ?? false;
+      const buttonStyle = collapsed ? style.expandButton : style.collapseButton;
 
       Object.entries(buttonStyle).forEach(([key, value]) => {
         if (Object.hasOwn(collapseButton, key)) {
@@ -582,27 +528,32 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
       buttonBound.h = buttonStyle.height;
       collapseButton.xywh = buttonBound.serialize();
       collapseButton.groupId = this.id;
-      collapseButton.opacity = 0;
+      collapseButton.text = collapsed ? node.children.length.toString() : '';
     }
 
     if (!btnExisted) {
+      collapseButton.opacity = 0;
+
       this.collapseButtons.set(id, collapseButton);
       this.surface.addLocalElement(collapseButton);
 
       collapseButton.onPointerEnter = () => {
-        collapseButton.opacity = 1;
+        const latestNode = this.getNode(node.id);
+        if (latestNode && latestNode.children.length > 0) {
+          collapseButton.opacity = 1;
+        }
       };
       collapseButton.onPointerLeave = () => {
         collapseButton.opacity = 0;
       };
       collapseButton.onClick = () => {
-        const nodeDetail = this.children.get(node.id);
+        const latestNode = this.getNode(node.id);
 
-        if (nodeDetail) {
+        if (latestNode && latestNode.children.length > 0) {
           this.surface.doc.transact(() => {
             this.children.set(node.id, {
-              ...nodeDetail,
-              collapsed: nodeDetail.collapsed ? false : true,
+              ...latestNode.detail,
+              collapsed: latestNode.detail.collapsed ? false : true,
             });
           });
         }
@@ -617,12 +568,44 @@ export class MindmapElementModel extends GfxGroupLikeElementModel<MindmapElement
       return null;
     }
 
-    return this.addConnector(
+    const layout = this.getLayoutDir(to)!;
+    const id = `#${from.id}-${to.id}`;
+    const connectorExist = this.connectors.has(id);
+    const connectorStyle = this.styleGetter.getNodeStyle(
+      to,
+      this.getPath(to)
+    ).connector;
+    const connector =
+      this.connectors.get(id) ?? new LocalConnectorElementModel(this.surface);
+    const outdated = this._isConnectorOutdated({
+      connector,
       from,
       to,
-      this.getLayoutDir(to)!,
-      this.styleGetter.getNodeStyle(to, this.getPath(to)).connector
-    );
+      layout,
+    });
+
+    if (!connectorExist) {
+      this.connectors.set(id, connector);
+    }
+
+    if (outdated) {
+      connector.id = id;
+      connector.source = {
+        id: from.id,
+        position: layout === LayoutType.RIGHT ? [1, 0.5] : [0, 0.5],
+      };
+      connector.target = {
+        id: to.id,
+        position: layout === LayoutType.RIGHT ? [0, 0.5] : [1, 0.5],
+      };
+
+      Object.entries(connectorStyle).forEach(([key, value]) => {
+        // @ts-ignore
+        connector[key as unknown] = value;
+      });
+    }
+
+    return { outdated, connector };
   }
 
   getLayoutDir(node: string | MindmapNode): LayoutType {
diff --git a/packages/affine/model/src/elements/mindmap/style.ts b/packages/affine/model/src/elements/mindmap/style.ts
index 351820d682259..fd7e4dface5c6 100644
--- a/packages/affine/model/src/elements/mindmap/style.ts
+++ b/packages/affine/model/src/elements/mindmap/style.ts
@@ -137,7 +137,7 @@ export class StyleOne extends MindmapStyleGetter {
         fillColor: '--affine-white',
 
         strokeColor: color,
-        strokeWidth: 2,
+        strokeWidth: 3,
       },
       expandButton: {
         width: 24,
@@ -148,7 +148,9 @@ export class StyleOne extends MindmapStyleGetter {
         fillColor: color,
 
         strokeColor: color,
-        strokeWidth: 2,
+        strokeWidth: 0,
+
+        padding: [4, 0],
 
         color: '--affine-white',
 
diff --git a/packages/affine/model/src/elements/shape/shape.ts b/packages/affine/model/src/elements/shape/shape.ts
index b20b6687a278b..49810dcf8f97c 100644
--- a/packages/affine/model/src/elements/shape/shape.ts
+++ b/packages/affine/model/src/elements/shape/shape.ts
@@ -199,7 +199,7 @@ export class LocalShapeElementModel extends GfxLocalElementModel {
 
   fontFamily: string = FontFamily.Inter;
 
-  fontSize!: number;
+  fontSize: number = 16;
 
   fontStyle: FontStyle = FontStyle.Normal;