summaryrefslogtreecommitdiffstats
path: root/vendor/ena/src/unify/mod.rs
blob: a26d699d8fd29dbcca72561e1cfc5541781026dd (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Union-find implementation. The main type is `UnificationTable`.
//!
//! You can define your own type for the *keys* in the table, but you
//! must implement `UnifyKey` for that type. The assumption is that
//! keys will be newtyped integers, hence we require that they
//! implement `Copy`.
//!
//! Keys can have values associated with them. The assumption is that
//! these values are cheaply cloneable (ideally, `Copy`), and some of
//! the interfaces are oriented around that assumption. If you just
//! want the classical "union-find" algorithm where you group things
//! into sets, use the `Value` type of `()`.
//!
//! When you have keys with non-trivial values, you must also define
//! how those values can be merged. As part of doing this, you can
//! define the "error" type to return on error; if errors are not
//! possible, use `NoError` (an uninstantiable struct). Using this
//! type also unlocks various more ergonomic methods (e.g., `union()`
//! in place of `unify_var_var()`).
//!
//! The best way to see how it is used is to read the `tests.rs` file;
//! search for e.g. `UnitKey`.

use std::fmt::Debug;
use std::marker;
use std::ops::Range;

use snapshot_vec::{self as sv, UndoLog};
use undo_log::{UndoLogs, VecLog};

mod backing_vec;
pub use self::backing_vec::{
    Delegate, InPlace, UnificationStore, UnificationStoreBase, UnificationStoreMut,
};

#[cfg(feature = "persistent")]
pub use self::backing_vec::Persistent;

#[cfg(test)]
mod tests;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Clients are expected to provide implementations of this trait; you
/// can see some examples in the `test` module.
pub trait UnifyKey: Copy + Clone + Debug + PartialEq {
    type Value: UnifyValue;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag() -> &'static str;

    /// If true, then `self` should be preferred as root to `other`.
    /// Note that we assume a consistent partial ordering, so
    /// returning true implies that `other.prefer_as_root_to(self)`
    /// would return false.  If there is no ordering between two keys
    /// (i.e., `a.prefer_as_root_to(b)` and `b.prefer_as_root_to(a)`
    /// both return false) then the rank will be used to determine the
    /// root in an optimal way.
    ///
    /// NB. The only reason to implement this method is if you want to
    /// control what value is returned from `find()`. In general, it
    /// is better to let the unification table determine the root,
    /// since overriding the rank can cause execution time to increase
    /// dramatically.
    #[allow(unused_variables)]
    fn order_roots(
        a: Self,
        a_value: &Self::Value,
        b: Self,
        b_value: &Self::Value,
    ) -> Option<(Self, Self)> {
        None
    }
}

/// Trait implemented for **values** associated with a unification
/// key. This trait defines how to merge the values from two keys that
/// are unioned together. This merging can be fallible. If you attempt
/// to union two keys whose values cannot be merged, then the error is
/// propagated up and the two keys are not unioned.
///
/// This crate provides implementations of `UnifyValue` for `()`
/// (which is infallible) and `Option<T>` (where `T: UnifyValue`). The
/// option implementation merges two sum-values using the `UnifyValue`
/// implementation of `T`.
///
/// See also `EqUnifyValue`, which is a convenience trait for cases
/// where the "merge" operation succeeds only if the two values are
/// equal.
pub trait UnifyValue: Clone + Debug {
    /// Defines the type to return when merging of two values fails.
    /// If merging is infallible, use the special struct `NoError`
    /// found in this crate, which unlocks various more convenient
    /// methods on the unification table.
    type Error;

    /// Given two values, produce a new value that combines them.
    /// If that is not possible, produce an error.
    fn unify_values(value1: &Self, value2: &Self) -> Result<Self, Self::Error>;
}

/// A convenient helper for unification values which must be equal or
/// else an error occurs. For example, if you are unifying types in a
/// simple functional language, this may be appropriate, since (e.g.)
/// you can't unify a type variable bound to `int` with one bound to
/// `float` (but you can unify two type variables both bound to
/// `int`).
///
/// Any type which implements `EqUnifyValue` automatially implements
/// `UnifyValue`; if the two values are equal, merging is permitted.
/// Otherwise, the error `(v1, v2)` is returned, where `v1` and `v2`
/// are the two unequal values.
pub trait EqUnifyValue: Eq + Clone + Debug {}

impl<T: EqUnifyValue> UnifyValue for T {
    type Error = (T, T);

    fn unify_values(value1: &Self, value2: &Self) -> Result<Self, Self::Error> {
        if value1 == value2 {
            Ok(value1.clone())
        } else {
            Err((value1.clone(), value2.clone()))
        }
    }
}

/// A struct which can never be instantiated. Used
/// for the error type for infallible cases.
#[derive(Debug)]
pub struct NoError {
    _dummy: (),
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq, Clone, Debug)]
pub struct VarValue<K: UnifyKey> {
    parent: K,       // if equal to self, this is a root
    value: K::Value, // value assigned (only relevant to root)
    rank: u32,       // max depth (only relevant to root)
}

/// Table of unification keys and their values. You must define a key type K
/// that implements the `UnifyKey` trait. Unification tables can be used in two-modes:
///
/// - in-place (`UnificationTable<InPlace<K>>` or `InPlaceUnificationTable<K>`):
///   - This is the standard mutable mode, where the array is modified
///     in place.
///   - To do backtracking, you can employ the `snapshot` and `rollback_to`
///     methods.
/// - persistent (`UnificationTable<Persistent<K>>` or `PersistentUnificationTable<K>`):
///   - In this mode, we use a persistent vector to store the data, so that
///     cloning the table is an O(1) operation.
///   - This implies that ordinary operations are quite a bit slower though.
///   - Requires the `persistent` feature be selected in your Cargo.toml file.
#[derive(Clone, Debug, Default)]
pub struct UnificationTable<S: UnificationStoreBase> {
    /// Indicates the current value of each key.
    values: S,
}

pub type UnificationStorage<K> = Vec<VarValue<K>>;
pub type UnificationTableStorage<K> = UnificationTable<InPlace<K, UnificationStorage<K>, ()>>;

/// A unification table that uses an "in-place" vector.
#[allow(type_alias_bounds)]
pub type InPlaceUnificationTable<
    K: UnifyKey,
    V: sv::VecLike<Delegate<K>> = Vec<VarValue<K>>,
    L = VecLog<UndoLog<Delegate<K>>>,
> = UnificationTable<InPlace<K, V, L>>;

/// A unification table that uses a "persistent" vector.
#[cfg(feature = "persistent")]
#[allow(type_alias_bounds)]
pub type PersistentUnificationTable<K: UnifyKey> = UnificationTable<Persistent<K>>;

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<S: UnificationStore> {
    // Link snapshot to the unification store `S` of the table.
    marker: marker::PhantomData<S>,
    snapshot: S::Snapshot,
}

impl<K: UnifyKey> VarValue<K> {
    fn new_var(key: K, value: K::Value) -> VarValue<K> {
        VarValue::new(key, value, 0)
    }

    fn new(parent: K, value: K::Value, rank: u32) -> VarValue<K> {
        VarValue {
            parent: parent, // this is a root
            value: value,
            rank: rank,
        }
    }

    fn redirect(&mut self, to: K) {
        self.parent = to;
    }

    fn root(&mut self, rank: u32, value: K::Value) {
        self.rank = rank;
        self.value = value;
    }

    fn parent(&self, self_key: K) -> Option<K> {
        self.if_not_self(self.parent, self_key)
    }

    fn if_not_self(&self, key: K, self_key: K) -> Option<K> {
        if key == self_key {
            None
        } else {
            Some(key)
        }
    }
}
impl<K> UnificationTableStorage<K>
where
    K: UnifyKey,
{
    /// Creates a `UnificationTable` using an external `undo_log`, allowing mutating methods to be
    /// called if `L` does not implement `UndoLogs`
    pub fn with_log<'a, L>(
        &'a mut self,
        undo_log: L,
    ) -> UnificationTable<InPlace<K, &'a mut UnificationStorage<K>, L>>
    where
        L: UndoLogs<sv::UndoLog<Delegate<K>>>,
    {
        UnificationTable {
            values: InPlace {
                values: self.values.values.with_log(undo_log),
            },
        }
    }
}

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<S: UnificationStoreBase + Default> UnificationTable<S> {
    pub fn new() -> Self {
        Self::default()
    }
}

impl<S: UnificationStore> UnificationTable<S> {
    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<S> {
        Snapshot {
            marker: marker::PhantomData::<S>,
            snapshot: self.values.start_snapshot(),
        }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<S>) {
        debug!("{}: rollback_to()", S::tag());
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<S>) {
        debug!("{}: commit()", S::tag());
        self.values.commit(snapshot.snapshot);
    }

    /// Returns the keys of all variables created since the `snapshot`.
    pub fn vars_since_snapshot(&self, snapshot: &Snapshot<S>) -> Range<S::Key> {
        let range = self.values.values_since_snapshot(&snapshot.snapshot);
        S::Key::from_index(range.start as u32)..S::Key::from_index(range.end as u32)
    }
}

impl<S: UnificationStoreBase> UnificationTable<S> {
    /// Returns the number of keys created so far.
    pub fn len(&self) -> usize {
        self.values.len()
    }
}

impl<S: UnificationStoreMut> UnificationTable<S> {
    /// Starts a new snapshot. Each snapshot must be either
    /// Creates a fresh key with the given value.
    pub fn new_key(&mut self, value: S::Value) -> S::Key {
        let len = self.values.len();
        let key: S::Key = UnifyKey::from_index(len as u32);
        self.values.push(VarValue::new_var(key, value));
        debug!("{}: created new key: {:?}", S::tag(), key);
        key
    }

    /// Reserve memory for `num_new_keys` to be created. Does not
    /// actually create the new keys; you must then invoke `new_key`.
    pub fn reserve(&mut self, num_new_keys: usize) {
        self.values.reserve(num_new_keys);
    }

    /// Clears all unifications that have been performed, resetting to
    /// the initial state. The values of each variable are given by
    /// the closure.
    pub fn reset_unifications(&mut self, mut value: impl FnMut(S::Key) -> S::Value) {
        self.values.reset_unifications(|i| {
            let key = UnifyKey::from_index(i as u32);
            let value = value(key);
            VarValue::new_var(key, value)
        });
    }

    /// Obtains the current value for a particular key.
    /// Not for end-users; they can use `probe_value`.
    fn value(&self, key: S::Key) -> &VarValue<S::Key> {
        &self.values[key.index() as usize]
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    ///
    /// This is an always-inlined version of this function for the hot
    /// callsites. `uninlined_get_root_key` is the never-inlined version.
    #[inline(always)]
    fn inlined_get_root_key(&mut self, vid: S::Key) -> S::Key {
        let redirect = {
            match self.value(vid).parent(vid) {
                None => return vid,
                Some(redirect) => redirect,
            }
        };

        let root_key: S::Key = self.uninlined_get_root_key(redirect);
        if root_key != redirect {
            // Path compression
            self.update_value(vid, |value| value.parent = root_key);
        }

        root_key
    }

    // This is a never-inlined version of this function for cold callsites.
    // 'inlined_get_root_key` is the always-inlined version.
    #[inline(never)]
    fn uninlined_get_root_key(&mut self, vid: S::Key) -> S::Key {
        self.inlined_get_root_key(vid)
    }

    fn update_value<OP>(&mut self, key: S::Key, op: OP)
    where
        OP: FnOnce(&mut VarValue<S::Key>),
    {
        self.values.update(key.index() as usize, op);
        debug!("Updated variable {:?} to {:?}", key, self.value(key));
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify_roots(&mut self, key_a: S::Key, key_b: S::Key, new_value: S::Value) {
        debug!("unify(key_a={:?}, key_b={:?})", key_a, key_b);

        let rank_a = self.value(key_a).rank;
        let rank_b = self.value(key_b).rank;
        if let Some((new_root, redirected)) = S::Key::order_roots(
            key_a,
            &self.value(key_a).value,
            key_b,
            &self.value(key_b).value,
        ) {
            // compute the new rank for the new root that they chose;
            // this may not be the optimal choice.
            let new_rank = if new_root == key_a {
                debug_assert!(redirected == key_b);
                if rank_a > rank_b {
                    rank_a
                } else {
                    rank_b + 1
                }
            } else {
                debug_assert!(new_root == key_b);
                debug_assert!(redirected == key_a);
                if rank_b > rank_a {
                    rank_b
                } else {
                    rank_a + 1
                }
            };
            self.redirect_root(new_rank, redirected, new_root, new_value);
        } else if rank_a > rank_b {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.redirect_root(rank_a, key_b, key_a, new_value);
        } else if rank_a < rank_b {
            // b has greater rank, so a should redirect to b.
            self.redirect_root(rank_b, key_a, key_b, new_value);
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            self.redirect_root(rank_a + 1, key_a, key_b, new_value);
        }
    }

    /// Internal method to redirect `old_root_key` (which is currently
    /// a root) to a child of `new_root_key` (which will remain a
    /// root). The rank and value of `new_root_key` will be updated to
    /// `new_rank` and `new_value` respectively.
    fn redirect_root(
        &mut self,
        new_rank: u32,
        old_root_key: S::Key,
        new_root_key: S::Key,
        new_value: S::Value,
    ) {
        self.update_value(old_root_key, |old_root_value| {
            old_root_value.redirect(new_root_key);
        });
        self.update_value(new_root_key, |new_root_value| {
            new_root_value.root(new_rank, new_value);
        });
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// Public API

impl<S, K, V> UnificationTable<S>
where
    S: UnificationStoreMut<Key = K, Value = V>,
    K: UnifyKey<Value = V>,
    V: UnifyValue,
{
    /// Unions two keys without the possibility of failure; only
    /// applicable when unify values use `NoError` as their error
    /// type.
    pub fn union<K1, K2>(&mut self, a_id: K1, b_id: K2)
    where
        K1: Into<K>,
        K2: Into<K>,
        V: UnifyValue<Error = NoError>,
    {
        self.unify_var_var(a_id, b_id).unwrap();
    }

    /// Unions a key and a value without the possibility of failure;
    /// only applicable when unify values use `NoError` as their error
    /// type.
    pub fn union_value<K1>(&mut self, id: K1, value: V)
    where
        K1: Into<K>,
        V: UnifyValue<Error = NoError>,
    {
        self.unify_var_value(id, value).unwrap();
    }

    /// Given two keys, indicates whether they have been unioned together.
    pub fn unioned<K1, K2>(&mut self, a_id: K1, b_id: K2) -> bool
    where
        K1: Into<K>,
        K2: Into<K>,
    {
        self.find(a_id) == self.find(b_id)
    }

    /// Given a key, returns the (current) root key.
    pub fn find<K1>(&mut self, id: K1) -> K
    where
        K1: Into<K>,
    {
        let id = id.into();
        self.uninlined_get_root_key(id)
    }

    /// Unions together two variables, merging their values. If
    /// merging the values fails, the error is propagated and this
    /// method has no effect.
    pub fn unify_var_var<K1, K2>(&mut self, a_id: K1, b_id: K2) -> Result<(), V::Error>
    where
        K1: Into<K>,
        K2: Into<K>,
    {
        let a_id = a_id.into();
        let b_id = b_id.into();

        let root_a = self.uninlined_get_root_key(a_id);
        let root_b = self.uninlined_get_root_key(b_id);

        if root_a == root_b {
            return Ok(());
        }

        let combined = V::unify_values(&self.value(root_a).value, &self.value(root_b).value)?;

        Ok(self.unify_roots(root_a, root_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`, attempting to merge
    /// with the previous value.
    pub fn unify_var_value<K1>(&mut self, a_id: K1, b: V) -> Result<(), V::Error>
    where
        K1: Into<K>,
    {
        let a_id = a_id.into();
        let root_a = self.uninlined_get_root_key(a_id);
        let value = V::unify_values(&self.value(root_a).value, &b)?;
        self.update_value(root_a, |node| node.value = value);
        Ok(())
    }

    /// Returns the current value for the given key. If the key has
    /// been union'd, this will give the value from the current root.
    pub fn probe_value<K1>(&mut self, id: K1) -> V
    where
        K1: Into<K>,
    {
        self.inlined_probe_value(id)
    }

    // An always-inlined version of `probe_value`, for hot callsites.
    #[inline(always)]
    pub fn inlined_probe_value<K1>(&mut self, id: K1) -> V
    where
        K1: Into<K>,
    {
        let id = id.into();
        let id = self.inlined_get_root_key(id);
        self.value(id).value.clone()
    }
}

///////////////////////////////////////////////////////////////////////////

impl UnifyValue for () {
    type Error = NoError;

    fn unify_values(_: &(), _: &()) -> Result<(), NoError> {
        Ok(())
    }
}

impl<V: UnifyValue> UnifyValue for Option<V> {
    type Error = V::Error;

    fn unify_values(a: &Option<V>, b: &Option<V>) -> Result<Self, V::Error> {
        match (a, b) {
            (&None, &None) => Ok(None),
            (&Some(ref v), &None) | (&None, &Some(ref v)) => Ok(Some(v.clone())),
            (&Some(ref a), &Some(ref b)) => match V::unify_values(a, b) {
                Ok(v) => Ok(Some(v)),
                Err(err) => Err(err),
            },
        }
    }
}