mol-data: fix SortedArray.union & mol-model: fix includeConnected

src/mol-data/int/_spec/sorted-array.spec.ts

@@ -134,4 +134,14 @@ describe('sortedArray', () => {
             SortedArray.ofSortedArray([768, 769, 770, 771, 772, 773, 774, 775, 776, 819, 820, 821, 822, 823, 824, 825, 826, 827, 828, 829, 1751, 1752, 1753, 1754, 1755, 1756, 1757, 1758, 1811, 1812, 1813, 1814, 1815, 1816, 1817, 1818, 1819])
         )
     })
+
+    it('union7', () => {
+        compareArrays(
+            SortedArray.union(
+                SortedArray.ofSortedArray([3766, 3767, 3768, 3770, 3773, 3780, 3783, 3787, 3797]),
+                SortedArray.ofSortedArray([3769, 3790, 3794])
+            ),
+            SortedArray.ofSortedArray([3766, 3767, 3768, 3769, 3770, 3773, 3780, 3783, 3787, 3790, 3794, 3797])
+        )
+    })
 });

src/mol-data/int/impl/sorted-array.ts

@@ -202,6 +202,10 @@ export function union(a: Nums, b: Nums): Nums {
         else { indices[offset++] = x; i++; j++; }
     }
 
+    // insert the remaining common part
+    for (; i < endI; i++) indices[offset++] = a[i];
+    for (; j < endJ; j++) indices[offset++] = b[j];
+
     // insert the "tail"
     for (; i < lenA; i++) indices[offset++] = a[i];
     for (; j < lenB; j++) indices[offset++] = b[j];

src/mol-model/structure/query/queries/modifiers.ts

@@ -320,6 +320,7 @@ export function includeConnected({ query, layerCount, wholeResidues, linkTest }:
             if (sI % 10 === 0) ctx.throwIfTimedOut();
         });
         ctx.popCurrentLink();
+
         return builder.getSelection();
     }
 }
@@ -335,7 +336,7 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
     const interLinks = inputStructure.links;
     const builder = new StructureUniqueSubsetBuilder(inputStructure);
 
-    const processedUnits = new Set<number>();
+    // Note: each link is visited twice so that link.atom-a and link.atom-b both get the "swapped values"
 
     const atomicLink = ctx.atomicLink;
 
@@ -380,8 +381,6 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
 
         // Process inter unit links
         for (const linkedUnit of interLinks.getLinkedUnits(inputUnit)) {
-            if (processedUnits.has(linkedUnit.unitA.id)) continue;
-
             atomicLink.b.unit = linkedUnit.unitB;
             for (const aI of linkedUnit.linkedElementIndices) {
                 // check if the element is in the expanded structure
@@ -401,8 +400,6 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
                 }
             }
         }
-
-        processedUnits.add(unit.id);
     }
 
     return builder.getStructure();