Struct cranelift_egraph::NodeKey

source ·

pub struct NodeKey { /* private fields */ }

Expand description

A reference to a node.

Implementations§

source §

impl NodeKey

source

pub fn node<'a, N>(&self, nodes: &'a [N]) -> &'a N

Get the node for this NodeKey, given the nodes from the appropriate EGraph.

Examples found in repository ?

src/lib.rs (line 339)

    fn ctx_eq(&self, a: &NodeKey, b: &NodeKey, uf: &mut UnionFind) -> bool {
        let a = a.node(self.nodes);
        let b = b.node(self.nodes);
        self.node_ctx.ctx_eq(a, b, uf)
    }
}

impl<'a, 'b, L: Language> CtxHash<NodeKey> for NodeKeyCtx<'a, 'b, L> {
    fn ctx_hash(&self, value: &NodeKey, uf: &mut UnionFind) -> u64 {
        self.node_ctx.ctx_hash(value.node(self.nodes), uf)
    }
}

/// An EClass entry. Contains either a single new enode and a child
/// eclass (i.e., adds one new enode), or unions two child eclasses
/// together.
#[derive(Debug, Clone, Copy)]
pub struct EClass {
    // formats:
    //
    // 00 | unused  (31 bits)        | NodeKey (31 bits)
    // 01 | eclass_child   (31 bits) | NodeKey (31 bits)
    // 10 | eclass_child_1 (31 bits) | eclass_child_id_2 (31 bits)
    bits: u64,
}

impl EClass {
    fn node(node: NodeKey) -> EClass {
        let node_idx = node.bits() as u64;
        debug_assert!(node_idx < (1 << 31));
        EClass {
            bits: (0b00 << 62) | node_idx,
        }
    }

    fn node_and_child(node: NodeKey, eclass_child: Id) -> EClass {
        let node_idx = node.bits() as u64;
        debug_assert!(node_idx < (1 << 31));
        debug_assert!(eclass_child != Id::invalid());
        let child = eclass_child.0 as u64;
        debug_assert!(child < (1 << 31));
        EClass {
            bits: (0b01 << 62) | (child << 31) | node_idx,
        }
    }

    fn union(child1: Id, child2: Id) -> EClass {
        debug_assert!(child1 != Id::invalid());
        let child1 = child1.0 as u64;
        debug_assert!(child1 < (1 << 31));

        debug_assert!(child2 != Id::invalid());
        let child2 = child2.0 as u64;
        debug_assert!(child2 < (1 << 31));

        EClass {
            bits: (0b10 << 62) | (child1 << 31) | child2,
        }
    }

    /// Get the node, if any, from a node-only or node-and-child
    /// eclass.
    pub fn get_node(&self) -> Option<NodeKey> {
        self.as_node()
            .or_else(|| self.as_node_and_child().map(|(node, _)| node))
    }

    /// Get the first child, if any.
    pub fn child1(&self) -> Option<Id> {
        self.as_node_and_child()
            .map(|(_, p1)| p1)
            .or(self.as_union().map(|(p1, _)| p1))
    }

    /// Get the second child, if any.
    pub fn child2(&self) -> Option<Id> {
        self.as_union().map(|(_, p2)| p2)
    }

    /// If this EClass is just a lone enode, return it.
    pub fn as_node(&self) -> Option<NodeKey> {
        if (self.bits >> 62) == 0b00 {
            let node_idx = (self.bits & ((1 << 31) - 1)) as u32;
            Some(NodeKey::from_bits(node_idx))
        } else {
            None
        }
    }

    /// If this EClass is one new enode and a child, return the node
    /// and child ID.
    pub fn as_node_and_child(&self) -> Option<(NodeKey, Id)> {
        if (self.bits >> 62) == 0b01 {
            let node_idx = (self.bits & ((1 << 31) - 1)) as u32;
            let child = ((self.bits >> 31) & ((1 << 31) - 1)) as u32;
            Some((NodeKey::from_bits(node_idx), Id::from_bits(child)))
        } else {
            None
        }
    }

    /// If this EClass is the union variety, return the two child
    /// EClasses. Both are guaranteed not to be `Id::invalid()`.
    pub fn as_union(&self) -> Option<(Id, Id)> {
        if (self.bits >> 62) == 0b10 {
            let child1 = ((self.bits >> 31) & ((1 << 31) - 1)) as u32;
            let child2 = (self.bits & ((1 << 31) - 1)) as u32;
            Some((Id::from_bits(child1), Id::from_bits(child2)))
        } else {
            None
        }
    }
}

/// A new or existing `T` when adding to a deduplicated set or data
/// structure, like an egraph.
#[derive(Clone, Copy, Debug)]
pub enum NewOrExisting<T> {
    New(T),
    Existing(T),
}

impl<T> NewOrExisting<T> {
    /// Get the underlying value.
    pub fn get(self) -> T {
        match self {
            NewOrExisting::New(t) => t,
            NewOrExisting::Existing(t) => t,
        }
    }
}

impl<L: Language, A: Analysis<L = L>> EGraph<L, A>
where
    L::Node: 'static,
{
    /// Create a new aegraph.
    pub fn new(analysis: Option<A>) -> Self {
        let analysis = analysis.map(|a| (a, SecondaryMap::new()));
        Self {
            nodes: vec![],
            node_map: CtxHashMap::new(),
            classes: PrimaryMap::new(),
            unionfind: UnionFind::new(),
            analysis,
        }
    }

    /// Create a new aegraph with the given capacity.
    pub fn with_capacity(nodes: usize, analysis: Option<A>) -> Self {
        let analysis = analysis.map(|a| (a, SecondaryMap::with_capacity(nodes)));
        Self {
            nodes: Vec::with_capacity(nodes),
            node_map: CtxHashMap::with_capacity(nodes),
            classes: PrimaryMap::with_capacity(nodes),
            unionfind: UnionFind::with_capacity(nodes),
            analysis,
        }
    }

    /// Add a new node.
    pub fn add(&mut self, node: L::Node, node_ctx: &L) -> NewOrExisting<Id> {
        // Push the node. We can then build a NodeKey that refers to
        // it and look for an existing interned copy. If one exists,
        // we can pop the pushed node and return the existing Id.
        let node_idx = self.nodes.len();
        trace!("adding node: {:?}", node);
        let needs_dedup = node_ctx.needs_dedup(&node);
        self.nodes.push(node);

        let key = NodeKey::from_node_idx(node_idx);
        if needs_dedup {
            let ctx = NodeKeyCtx {
                nodes: &self.nodes[..],
                node_ctx,
            };

            match self.node_map.entry(key, &ctx, &mut self.unionfind) {
                Entry::Occupied(o) => {
                    let eclass_id = *o.get();
                    self.nodes.pop();
                    trace!(" -> existing id {}", eclass_id);
                    NewOrExisting::Existing(eclass_id)
                }
                Entry::Vacant(v) => {
                    // We're creating a new eclass now.
                    let eclass_id = self.classes.push(EClass::node(key));
                    trace!(" -> new node and eclass: {}", eclass_id);
                    self.unionfind.add(eclass_id);

                    // Add to interning map with a NodeKey referring to the eclass.
                    v.insert(eclass_id);

                    // Update analysis.
                    let node_ctx = ctx.node_ctx;
                    self.update_analysis_new(node_ctx, eclass_id, key);

                    NewOrExisting::New(eclass_id)
                }
            }
        } else {
            let eclass_id = self.classes.push(EClass::node(key));
            self.unionfind.add(eclass_id);
            NewOrExisting::New(eclass_id)
        }
    }

    /// Merge one eclass into another, maintaining the acyclic
    /// property (args must have lower eclass Ids than the eclass
    /// containing the node with those args). Returns the Id of the
    /// merged eclass.
    pub fn union(&mut self, ctx: &L, a: Id, b: Id) -> Id {
        assert_ne!(a, Id::invalid());
        assert_ne!(b, Id::invalid());
        let (a, b) = (std::cmp::max(a, b), std::cmp::min(a, b));
        trace!("union: id {} and id {}", a, b);
        if a == b {
            trace!(" -> no-op");
            return a;
        }

        self.unionfind.union(a, b);

        // If the younger eclass has no child, we can link it
        // directly and return that eclass. Otherwise, we create a new
        // union eclass.
        if let Some(node) = self.classes[a].as_node() {
            trace!(
                " -> id {} is one-node eclass; making into node-and-child with id {}",
                a,
                b
            );
            self.classes[a] = EClass::node_and_child(node, b);
            self.update_analysis_union(ctx, a, a, b);
            return a;
        }

        let u = self.classes.push(EClass::union(a, b));
        self.unionfind.add(u);
        self.unionfind.union(u, b);
        trace!(" -> union id {} and id {} into id {}", a, b, u);
        self.update_analysis_union(ctx, u, a, b);
        u
    }

    /// Get the canonical ID for an eclass. This may be an older
    /// generation, so will not be able to see all enodes in the
    /// eclass; but it will allow us to unambiguously refer to an
    /// eclass, even across merging.
    pub fn canonical_id_mut(&mut self, eclass: Id) -> Id {
        self.unionfind.find_and_update(eclass)
    }

    /// Get the canonical ID for an eclass. This may be an older
    /// generation, so will not be able to see all enodes in the
    /// eclass; but it will allow us to unambiguously refer to an
    /// eclass, even across merging.
    pub fn canonical_id(&self, eclass: Id) -> Id {
        self.unionfind.find(eclass)
    }

    /// Get the enodes for a given eclass.
    pub fn enodes(&self, eclass: Id) -> NodeIter<L, A> {
        NodeIter {
            stack: smallvec![eclass],
            _phantom1: PhantomData,
            _phantom2: PhantomData,
        }
    }

    /// Update analysis for a given eclass node (new-enode case).
    fn update_analysis_new(&mut self, ctx: &L, eclass: Id, node: NodeKey) {
        if let Some((analysis, state)) = self.analysis.as_mut() {
            let node = node.node(&self.nodes);
            state[eclass] = analysis.for_node(ctx, node, state);
        }
    }