🐛 fix(tab-groups): address review feedback from gnachman

Author: kud

sources/PseudoTerminal+Private.h

@@ -126,6 +126,8 @@ extern NSString *const TERMINAL_ARRANGEMENT_SIZE_LOCKED;
 // the right type.
 @property (nonatomic, readonly) __unsafe_unretained iTermRootTerminalView *contentView;
 
+@property (nonatomic, readonly) NSMutableArray *mutableTabGroups;
+
 - (void)returnTabBarToContentView;
 - (void)updateForTransparency:(NSWindow<PTYWindow> *)window;
 - (void)updateVariables;

sources/PseudoTerminal+TabGroups.swift

@@ -8,15 +8,6 @@ extension PseudoTerminal {
         return mutableTabGroups as! [iTermTabGroup]
     }
 
-    private var mutableTabGroups: NSMutableArray {
-        if let existing = objc_getAssociatedObject(self, &AssociatedKeys.tabGroups) as? NSMutableArray {
-            return existing
-        }
-        let array = NSMutableArray()
-        objc_setAssociatedObject(self, &AssociatedKeys.tabGroups, array, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
-        return array
-    }
-
     @objc @discardableResult
     func createTabGroup(name: String, color: NSColor, forTab tab: PTYTab) -> iTermTabGroup {
         let group = iTermTabGroup(name: name, color: color)
@@ -45,13 +36,31 @@ extension PseudoTerminal {
 
     @objc func deleteTabGroup(_ group: iTermTabGroup) {
         if group.isCollapsed {
-            toggleCollapseGroup(group)
+            unstashItems(for: group)
         }
         group.memberTabIDs.removeAll()
         mutableTabGroups.remove(group)
         updateTabGroupDecorations()
     }
 
+    private func unstashItems(for group: iTermTabGroup) {
+        guard let representativeID = group.memberTabIDs.first,
+              let representativeTab = tabs()?.first(where: { Int($0.uniqueId) == representativeID }),
+              let representativeItem = representativeTab.tabViewItem else { return }
+
+        let tabView = contentView.tabView
+        let baseIndex = tabView.indexOfTabViewItem(representativeItem)
+        guard baseIndex != NSNotFound else { return }
+
+        for (offset, item) in group.stashedTabViewItems.enumerated() {
+            let targetIndex = baseIndex + 1 + offset
+            let safeIndex = min(targetIndex, tabView.numberOfTabViewItems)
+            tabView.insertTabViewItem(item, at: safeIndex)
+        }
+        group.stashedTabViewItems.removeAll()
+        group.isCollapsed = false
+    }
+
     @objc func toggleCollapseGroup(_ group: iTermTabGroup) {
         if group.isCollapsed {
             expandGroup(group)
@@ -104,18 +113,20 @@ extension PseudoTerminal {
     }
 
     @objc func encodeTabGroupsForArrangement() -> [[String: Any]] {
-        return tabGroups.map { $0.toDictionary() }
+        let allTabs = tabs() ?? []
+        return tabGroups.map { $0.toDictionary(allTabs: allTabs) }
     }
 
     @objc func restoreTabGroupsFromArrangement(_ arrangement: [[String: Any]]) {
-        let allTabs = tabs()
+        let allTabs = tabs() ?? []
         for dict in arrangement {
             guard let group = iTermTabGroup(dictionary: dict) else { continue }
-            let validIDs = group.memberTabIDs.filter { id in
-                allTabs.contains { Int($0.uniqueId) == id }
+            let resolvedIDs = group.memberTabIDs.compactMap { index -> Int? in
+                guard index >= 0, index < allTabs.count else { return nil }
+                return Int(allTabs[index].uniqueId)
             }
-            guard !validIDs.isEmpty else { continue }
-            group.memberTabIDs = validIDs
+            guard !resolvedIDs.isEmpty else { continue }
+            group.memberTabIDs = resolvedIDs
             mutableTabGroups.add(group)
         }
         updateTabGroupDecorations()
@@ -168,7 +179,7 @@ extension PseudoTerminal {
     }
 
     private func collapseGroup(_ group: iTermTabGroup) {
-        let allTabs = tabs()
+        let allTabs = tabs() ?? []
         let memberTabs = group.memberTabIDs.compactMap { id in allTabs.first { Int($0.uniqueId) == id } }
         guard let representative = memberTabs.first else { return }
 
@@ -189,7 +200,7 @@ extension PseudoTerminal {
 
     private func expandGroup(_ group: iTermTabGroup) {
         guard let representativeID = group.memberTabIDs.first,
-              let representativeTab = tabs().first(where: { Int($0.uniqueId) == representativeID }),
+              let representativeTab = (tabs() ?? []).first(where: { Int($0.uniqueId) == representativeID }),
               let representativeItem = representativeTab.tabViewItem else { return }
 
         let tabView = contentView.tabView
@@ -208,7 +219,7 @@ extension PseudoTerminal {
 
     private func enforceContiguity(of group: iTermTabGroup) {
         guard group.memberTabIDs.count > 1 else { return }
-        let allTabs = tabs()
+        let allTabs = tabs() ?? []
         let memberTabs = group.memberTabIDs.compactMap { id in allTabs.first { Int($0.uniqueId) == id } }
         guard let representativeItem = memberTabs.first?.tabViewItem else { return }
 
@@ -228,7 +239,3 @@ extension PseudoTerminal {
         }
     }
 }
-
-private enum AssociatedKeys {
-    static var tabGroups = "iTermTabGroupsKey"
-}

sources/PseudoTerminal.m

@@ -432,6 +432,8 @@ @implementation PseudoTerminal {
     iTermIdempotentOperationJoiner *_rightExtraJoiner;
     BOOL _excursionPrevented;
 
+    NSMutableArray<iTermTabGroup *> *_mutableTabGroups;
+
 }
 
 @synthesize scope = _scope;
@@ -457,10 +459,15 @@ - (instancetype)initWithWindowNibName:(NSString *)windowNibName {
     if (self) {
         _automaticallySelectNewTabs = YES;
         self.autoCommandHistorySessionGuid = nil;
+        _mutableTabGroups = [[NSMutableArray alloc] init];
     }
     return self;
 }
 
+- (NSMutableArray *)mutableTabGroups {
+    return _mutableTabGroups;
+}
+
 - (instancetype)initWithSmartLayout:(BOOL)smartLayout
                          windowType:(iTermWindowType)windowType
                     savedWindowType:(iTermWindowType)savedWindowType
@@ -1098,6 +1105,7 @@ - (void)dealloc {
     [_fullScreenEnteredSeal release];
     [_windowSizeHelper release];
     [_titlebarAccessoryNanny release];
+    [_mutableTabGroups release];
 
     [super dealloc];
 }

sources/iTermTabGroup.swift

@@ -13,6 +13,7 @@ import Cocoa
         static let name = "name"
         static let color = "color"
         static let memberTabIDs = "memberTabIDs"
+        static let memberTabIndices = "memberTabIndices"
         static let isCollapsed = "isCollapsed"
     }
 
@@ -54,11 +55,15 @@ import Cocoa
         coder.encode(isCollapsed, forKey: CodingKey.isCollapsed)
     }
 
-    @objc func toDictionary() -> [String: Any] {
+    @objc func toDictionary(allTabs: [PTYTab]) -> [String: Any] {
+        let indices = memberTabIDs.compactMap { id -> NSNumber? in
+            guard let index = allTabs.firstIndex(where: { Int($0.uniqueId) == id }) else { return nil }
+            return NSNumber(value: index)
+        }
         var dict: [String: Any] = [
             CodingKey.identifier: identifier,
             CodingKey.name: name,
-            CodingKey.memberTabIDs: memberTabIDs.map { NSNumber(value: $0) },
+            CodingKey.memberTabIndices: indices,
             CodingKey.isCollapsed: isCollapsed
         ]
         if let colorData = try? NSKeyedArchiver.archivedData(withRootObject: color, requiringSecureCoding: false) {
@@ -72,13 +77,13 @@ import Cocoa
               let name = dictionary[CodingKey.name] as? String,
               let colorData = dictionary[CodingKey.color] as? Data,
               let color = try? NSKeyedUnarchiver.unarchivedObject(ofClass: NSColor.self, from: colorData),
-              let rawIDs = dictionary[CodingKey.memberTabIDs] as? [NSNumber] else {
+              let rawIndices = dictionary[CodingKey.memberTabIndices] as? [NSNumber] else {
             return nil
         }
         self.identifier = identifier
         self.name = name
         self.color = color
-        self.memberTabIDs = rawIDs.map { $0.intValue }
+        self.memberTabIDs = rawIndices.map { $0.intValue }
         self.isCollapsed = (dictionary[CodingKey.isCollapsed] as? Bool) ?? false
     }
 }