1 files changed, 262 insertions, 149 deletions

diff --git a/apps/plugins/puzzles/src/dsf.c b/apps/plugins/puzzles/src/dsf.c
index 832bb3005a..bed564f8ec 100644
--- a/apps/plugins/puzzles/src/dsf.c
+++ b/apps/plugins/puzzles/src/dsf.c
@@ -5,188 +5,301 @@
  */
 
 #include <assert.h>
+#include <limits.h>
 #include <string.h>
 
 #include "puzzles.h"
 
-/*void print_dsf(int *dsf, int size)
+#define DSF_INDEX_MASK (UINT_MAX >> 1)
+#define DSF_FLAG_CANONICAL (UINT_MAX & ~(UINT_MAX >> 1))
+#define DSF_MAX (DSF_INDEX_MASK + 1)
+
+struct DSF {
+    /*
+     * Size of the dsf.
+     */
+    size_t size;
+
+    /*
+     * Main array storing the data structure.
+     *
+     * If n is the canonical element of an equivalence class,
+     * parent_or_size[n] holds the number of elements in that class,
+     * bitwise-ORed with DSF_FLAG_CANONICAL.
+     *
+     * If n is not the canonical element, parent_or_size[n] holds the
+     * index of another element nearer to the root of the tree for
+     * that class.
+     */
+    unsigned *parent_or_size;
+
+    /*
+     * Extra storage for flip tracking.
+     *
+     * If n is not a canonical element, flip[n] indicates whether the
+     * sense of this element is flipped relative to parent_or_size[n].
+     *
+     * If n is a canonical element, flip[n] is unused.
+     */
+    unsigned char *flip;
+
+    /*
+     * Extra storage for minimal-element tracking.
+     *
+     * If n is a canonical element, min[n] holds the index of the
+     * smallest value in n's equivalence class.
+     *
+     * If n is not a canonical element, min[n] is unused.
+     */
+    unsigned *min;
+};
+
+static DSF *dsf_new_internal(int size, bool flip, bool min)
 {
-    int *printed_elements = snewn(size, int);
-    int *equal_elements = snewn(size, int);
-    int *inverse_elements = snewn(size, int);
-    int printed_count = 0, equal_count, inverse_count;
-    int i, n;
-    bool inverse;
+    DSF *dsf;
 
-    memset(printed_elements, -1, sizeof(int) * size);
+    assert(0 < size && size <= DSF_MAX && "Bad dsf size");
 
-    while (1) {
-        equal_count = 0;
-        inverse_count = 0;
-        for (i = 0; i < size; ++i) {
-            if (!memchr(printed_elements, i, sizeof(int) * size)) 
-                break;
-        }
-        if (i == size)
-            goto done;
-
-        i = dsf_canonify(dsf, i);
-
-        for (n = 0; n < size; ++n) {
-            if (edsf_canonify(dsf, n, &inverse) == i) {
-               if (inverse)
-                   inverse_elements[inverse_count++] = n;
-               else
-                   equal_elements[equal_count++] = n;
-            }
-        }
-        
-        for (n = 0; n < equal_count; ++n) {
-            fprintf(stderr, "%d ", equal_elements[n]);
-            printed_elements[printed_count++] = equal_elements[n];
-        }
-        if (inverse_count) {
-            fprintf(stderr, "!= ");
-            for (n = 0; n < inverse_count; ++n) {
-                fprintf(stderr, "%d ", inverse_elements[n]);
-                printed_elements[printed_count++] = inverse_elements[n];
-            }
-        }
-        fprintf(stderr, "\n");
-    }
-done:
+    dsf = snew(DSF);
+    dsf->size = size;
+    dsf->parent_or_size = snewn(size, unsigned);
+    dsf->flip = flip ? snewn(size, unsigned char) : NULL;
+    dsf->min = min ? snewn(size, unsigned) : NULL;
+
+    dsf_reinit(dsf);
 
-    sfree(printed_elements);
-    sfree(equal_elements);
-    sfree(inverse_elements);
-}*/
+    return dsf;
+}
 
-void dsf_init(int *dsf, int size)
+DSF *dsf_new(int size)
 {
-    int i;
+    return dsf_new_internal(size, false, false);
+}
 
-    for (i = 0; i < size; i++) dsf[i] = 6;
-    /* Bottom bit of each element of this array stores whether that
-     * element is opposite to its parent, which starts off as
-     * false. Second bit of each element stores whether that element
-     * is the root of its tree or not.  If it's not the root, the
-     * remaining 30 bits are the parent, otherwise the remaining 30
-     * bits are the number of elements in the tree.  */
+DSF *dsf_new_flip(int size)
+{
+    return dsf_new_internal(size, true, false);
 }
 
-int *snew_dsf(int size)
+DSF *dsf_new_min(int size)
 {
-    int *ret;
-    
-    ret = snewn(size, int);
-    dsf_init(ret, size);
+    return dsf_new_internal(size, false, true);
+}
+
+void dsf_reinit(DSF *dsf)
+{
+    size_t i;
+
+    /* Every element starts as the root of an equivalence class of size 1 */
+    for (i = 0; i < dsf->size; i++)
+        dsf->parent_or_size[i] = DSF_FLAG_CANONICAL | 1;
 
-    /*print_dsf(ret, size); */
+    /* If we're tracking minima then every element is also its own min */
+    if (dsf->min)
+        for (i = 0; i < dsf->size; i++)
+            dsf->min[i] = i;
 
-    return ret;
+    /* No need to initialise dsf->flip, even if it exists, because
+     * only the entries for non-root elements are meaningful, and
+     * currently there are none. */
 }
 
-int dsf_canonify(int *dsf, int index)
+void dsf_copy(DSF *to, DSF *from)
 {
-    return edsf_canonify(dsf, index, NULL);
+    assert(to->size == from->size && "Mismatch in dsf_copy");
+    memcpy(to->parent_or_size, from->parent_or_size,
+           to->size * sizeof(*to->parent_or_size));
+    if (to->flip) {
+        assert(from->flip && "Copying a non-flip dsf to a flip one");
+        memcpy(to->flip, from->flip, to->size * sizeof(*to->flip));
+    }
+    if (to->min) {
+        assert(from->min && "Copying a non-min dsf to a min one");
+        memcpy(to->min, from->min, to->size * sizeof(*to->min));
+    }
 }
 
-void dsf_merge(int *dsf, int v1, int v2)
+
+void dsf_free(DSF *dsf)
 {
-    edsf_merge(dsf, v1, v2, false);
+    if (dsf) {
+        sfree(dsf->parent_or_size);
+        sfree(dsf->flip);
+        sfree(dsf->min);
+        sfree(dsf);
+    }
 }
 
-int dsf_size(int *dsf, int index) {
-    return dsf[dsf_canonify(dsf, index)] >> 2;
+static inline size_t dsf_find_root(DSF *dsf, size_t n)
+{
+    while (!(dsf->parent_or_size[n] & DSF_FLAG_CANONICAL))
+        n = dsf->parent_or_size[n];
+    return n;
+}
+
+static inline void dsf_path_compress(DSF *dsf, size_t n, size_t root)
+{
+    while (!(dsf->parent_or_size[n] & DSF_FLAG_CANONICAL)) {
+        size_t prev = n;
+        n = dsf->parent_or_size[n];
+        dsf->parent_or_size[prev] = root;
+    }
+    assert(n == root);
 }
 
-int edsf_canonify(int *dsf, int index, bool *inverse_return)
+int dsf_canonify(DSF *dsf, int n)
 {
-    int start_index = index, canonical_index;
-    bool inverse = false;
+    size_t root;
 
-/*    fprintf(stderr, "dsf = %p\n", dsf); */
-/*    fprintf(stderr, "Canonify %2d\n", index); */
+    assert(0 <= n && n < dsf->size && "Overrun in dsf_canonify");
 
-    assert(index >= 0);
+    root = dsf_find_root(dsf, n);
+    dsf_path_compress(dsf, n, root);
+    return root;
+}
+
+void dsf_merge(DSF *dsf, int n1, int n2)
+{
+    size_t r1, r2, s1, s2, root;
+
+    assert(0 <= n1 && n1 < dsf->size && "Overrun in dsf_merge");
+    assert(0 <= n2 && n2 < dsf->size && "Overrun in dsf_merge");
+    assert(!dsf->flip && "dsf_merge on a flip dsf");
+
+    /* Find the root elements */
+    r1 = dsf_find_root(dsf, n1);
+    r2 = dsf_find_root(dsf, n2);
+
+    if (r1 == r2) {
+        /* Classes are already the same, so we have a common root */
+        root = r1;
+    } else {
+        /* Classes must be merged */
+
+        /* Decide which one to use as the overall root, based on size */
+        s1 = dsf->parent_or_size[r1] & DSF_INDEX_MASK;
+        s2 = dsf->parent_or_size[r2] & DSF_INDEX_MASK;
+        if (s1 > s2) {
+            dsf->parent_or_size[r2] = root = r1;
+        } else {
+            dsf->parent_or_size[r1] = root = r2;
+        }
+        dsf->parent_or_size[root] = (s1 + s2) | DSF_FLAG_CANONICAL;
 
-    /* Find the index of the canonical element of the 'equivalence class' of
-     * which start_index is a member, and figure out whether start_index is the
-     * same as or inverse to that. */
-    while ((dsf[index] & 2) == 0) {
-        inverse ^= (dsf[index] & 1);
-        index = dsf[index] >> 2;
-/*        fprintf(stderr, "index = %2d, ", index); */
-/*        fprintf(stderr, "inverse = %d\n", inverse); */
+        if (dsf->min) {
+            /* Update the min of the merged class */
+            unsigned m1 = dsf->min[r1], m2 = dsf->min[r2];
+            dsf->min[root] = m1 < m2 ? m1 : m2;
+        }
     }
-    canonical_index = index;
-    
-    if (inverse_return)
-        *inverse_return = inverse;
-    
-    /* Update every member of this 'equivalence class' to point directly at the
-     * canonical member. */
-    index = start_index;
-    while (index != canonical_index) {
-        int nextindex = dsf[index] >> 2;
-        bool nextinverse = inverse ^ (dsf[index] & 1);
-        dsf[index] = (canonical_index << 2) | inverse;
-        inverse = nextinverse;
-        index = nextindex;
+
+    /* Path-compress both paths from n1 and n2 so they point at the new root */
+    dsf_path_compress(dsf, n1, root);
+    dsf_path_compress(dsf, n2, root);
+}
+
+bool dsf_equivalent(DSF *dsf, int n1, int n2)
+{
+    return dsf_canonify(dsf, n1) == dsf_canonify(dsf, n2);
+}
+
+int dsf_size(DSF *dsf, int n)
+{
+    size_t root = dsf_canonify(dsf, n);
+    return dsf->parent_or_size[root] & DSF_INDEX_MASK;
+}
+
+static inline size_t dsf_find_root_flip(DSF *dsf, size_t n, unsigned *flip)
+{
+    *flip = 0;
+    while (!(dsf->parent_or_size[n] & DSF_FLAG_CANONICAL)) {
+        *flip ^= dsf->flip[n];
+        n = dsf->parent_or_size[n];
+    }
+    return n;
+}
+
+static inline void dsf_path_compress_flip(DSF *dsf, size_t n, size_t root,
+                                          unsigned flip)
+{
+    while (!(dsf->parent_or_size[n] & DSF_FLAG_CANONICAL)) {
+        size_t prev = n;
+        unsigned flip_prev = flip;
+        n = dsf->parent_or_size[n];
+        flip ^= dsf->flip[prev];
+        dsf->flip[prev] = flip_prev;
+        dsf->parent_or_size[prev] = root;
     }
+    assert(n == root);
+}
+
+int dsf_canonify_flip(DSF *dsf, int n, bool *inverse)
+{
+    size_t root;
+    unsigned flip;
+
+    assert(0 <= n && n < dsf->size && "Overrun in dsf_canonify_flip");
+    assert(dsf->flip && "dsf_canonify_flip on a non-flip dsf");
+
+    root = dsf_find_root_flip(dsf, n, &flip);
+    dsf_path_compress_flip(dsf, n, root, flip);
+    *inverse = flip;
+    return root;
+}
+
+void dsf_merge_flip(DSF *dsf, int n1, int n2, bool inverse)
+{
+    size_t r1, r2, s1, s2, root;
+    unsigned f1, f2;
+
+    assert(0 <= n1 && n1 < dsf->size && "Overrun in dsf_merge_flip");
+    assert(0 <= n2 && n2 < dsf->size && "Overrun in dsf_merge_flip");
+    assert(dsf->flip && "dsf_merge_flip on a non-flip dsf");
+
+    /* Find the root elements */
+    r1 = dsf_find_root_flip(dsf, n1, &f1);
+    r2 = dsf_find_root_flip(dsf, n2, &f2);
+
+    if (r1 == r2) {
+        /* Classes are already the same, so we have a common root */
+        assert((f1 ^ f2 ^ inverse) == 0 && "Inconsistency in dsf_merge_flip");
+        root = r1;
+    } else {
+        /* Classes must be merged */
 
-    assert(!inverse);
-
-/*    fprintf(stderr, "Return %2d\n", index); */
-    
-    return index;
-}
-
-void edsf_merge(int *dsf, int v1, int v2, bool inverse)
-{
-    bool i1, i2;
-
-/*    fprintf(stderr, "dsf = %p\n", dsf); */
-/*    fprintf(stderr, "Merge [%2d,%2d], %d\n", v1, v2, inverse); */
-    
-    v1 = edsf_canonify(dsf, v1, &i1);
-    assert(dsf[v1] & 2);
-    inverse ^= i1;
-    v2 = edsf_canonify(dsf, v2, &i2);
-    assert(dsf[v2] & 2);
-    inverse ^= i2;
-
-/*    fprintf(stderr, "Doing [%2d,%2d], %d\n", v1, v2, inverse); */
-
-    if (v1 == v2)
-        assert(!inverse);
-    else {
-        /*
-         * We always make the smaller of v1 and v2 the new canonical
-         * element. This ensures that the canonical element of any
-         * class in this structure is always the first element in
-         * it. 'Keen' depends critically on this property.
-         *
-         * (Jonas Koelker previously had this code choosing which
-         * way round to connect the trees by examining the sizes of
-         * the classes being merged, so that the root of the
-         * larger-sized class became the new root. This gives better
-         * asymptotic performance, but I've changed it to do it this
-         * way because I like having a deterministic canonical
-         * element.)
-         */
-        if (v1 > v2) {
-            int v3 = v1;
-            v1 = v2;
-            v2 = v3;
-        }
-        dsf[v1] += (dsf[v2] >> 2) << 2;
-        dsf[v2] = (v1 << 2) | inverse;
+        /* Decide which one to use as the overall root, based on size */
+        s1 = dsf->parent_or_size[r1] & DSF_INDEX_MASK;
+        s2 = dsf->parent_or_size[r2] & DSF_INDEX_MASK;
+        if (s1 > s2) {
+            dsf->parent_or_size[r2] = root = r1;
+            dsf->flip[r2] = f1 ^ f2 ^ inverse;
+            f2 ^= dsf->flip[r2];
+        } else {
+            root = r2;
+            dsf->parent_or_size[r1] = root = r2;
+            dsf->flip[r1] = f1 ^ f2 ^ inverse;
+            f1 ^= dsf->flip[r1];
+        }
+        dsf->parent_or_size[root] = (s1 + s2) | DSF_FLAG_CANONICAL;
+
+        if (dsf->min) {
+            /* Update the min of the merged class */
+            unsigned m1 = dsf->min[r1], m2 = dsf->min[r2];
+            dsf->min[root] = m1 < m2 ? m1 : m2;
+        }
     }
-    
-    v2 = edsf_canonify(dsf, v2, &i2);
-    assert(v2 == v1);
-    assert(i2 == inverse);
 
-/*    fprintf(stderr, "dsf[%2d] = %2d\n", v2, dsf[v2]); */
+    /* Path-compress both paths from n1 and n2 so they point at the new root */
+    dsf_path_compress_flip(dsf, n1, root, f1);
+    dsf_path_compress_flip(dsf, n2, root, f2);
+}
+
+int dsf_minimal(DSF *dsf, int n)
+{
+    size_t root;
+
+    assert(dsf->min && "dsf_minimal on a non-min dsf");
+
+    root = dsf_canonify(dsf, n);
+    return dsf->min[root];
 }