libreda-logic 0.0.3

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Generic data-structure for logic networks. The node type is a type parameter.

use crate::linked_lists::{ElementIndex, LinkedLists, ListIndex};
use crate::network_simulator;
use crate::traits::{BooleanSystem, NumInputs, NumOutputs, PartialBooleanSystem};
use crate::truth_table::small_lut::truth_table_library;

use crate::network::*;
use crate::truth_table::small_lut::SmallTruthTable;
use itertools::Itertools;
use smallvec::SmallVec;

use fnv::FnvHashMap;
use std::hash::Hash;

use super::SimplifyResult;

/// An implementation of logic networks which is parameterized with the type of the logic node.
#[derive(Clone, Default)]
pub struct LogicNetwork<Node> {
    storage: Vec<NodeWithReferences<Node>>,
    primary_inputs: Vec<NodeId>,
    primary_outputs: Vec<NodeId>,
    /// Table for structural hashing.
    /// Structural hashing helps de-duplicating nodes.
    hash_table: FnvHashMap<Node, NodeId>,

    /// Storage for reverse graph edges.
    references: LinkedLists<NodeId>,
    edges: LinkedLists<(NodeId, ElementIndex)>,
}

impl<Node> LogicNetwork<Node> {
    /// Create a new and empty logic network.
    pub fn new() -> Self {
        Self {
            storage: Default::default(),
            primary_inputs: Default::default(),
            primary_outputs: Default::default(),
            hash_table: Default::default(),
            references: Default::default(),
            edges: Default::default(),
        }
    }
}

/// Associate a list of incoming edges with a network node.
#[derive(Clone, Default)]
struct NodeWithReferences<N> {
    /// The network node.
    node: N,
    /// Pointer to the start of a linked list. The linked list is stored in `LogicNetwork::references`.
    /// The list contains the indices of the nodes which use this node.
    references: ListIndex,
    outgoing_edges: ListIndex,
}

impl<N> From<N> for NodeWithReferences<N> {
    fn from(n: N) -> Self {
        Self {
            node: n,
            references: ListIndex::Empty,
            outgoing_edges: ListIndex::Empty,
        }
    }
}

/// ID of a signal in the logic network.
/// An ID can point to the output of a node as well as to a primary input or a constant.
/// The ID encodes if the signal is inverted.
#[derive(Clone, Copy, Hash, PartialEq, PartialOrd, Eq, Ord)]
pub struct NodeId {
    /// ID of the graph edge. This encodes the index into the storage array, an inversion bit and an 'input' bit.
    index: usize,
}

impl std::fmt::Debug for NodeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_constant() {
            write!(
                f,
                "NodeId(constant = {})",
                if self.is_zero() { "0" } else { "1" }
            )
        } else if self.is_input() {
            write!(
                f,
                "NodeId(input({}), inverted = {})",
                self.input_index().unwrap(),
                self.is_inverted()
            )
        } else {
            write!(
                f,
                "NodeId(index: {}, is_inverted: {})",
                self.node_index().unwrap(),
                self.is_inverted(),
            )
        }
    }
}

impl NodeId {
    /// Get the ID for the constant 0.
    pub(super) fn zero() -> Self {
        Self { index: 0 }
    }

    /// Get the index into the storage array.
    /// Return `None` if the ID points to a constant or an input.
    pub(super) fn node_index(&self) -> Option<usize> {
        (!self.is_constant() && !self.is_input()).then(|| {
            let masked = self.index >> 1; // Strip away the 'input' and 'inversion' bits.
            masked - 1
        })
    }

    /// Get the index of the primary input if this node ID represents a primary input.
    pub fn input_index(&self) -> Option<usize> {
        self.is_input().then(|| {
            let input_marker_bit = 1 << (usize::BITS - 1);
            (!input_marker_bit & self.index) >> 1 // Strip away the 'input' and 'inversion' bits.
        })
    }

    /// Create a new node ID which points to the node stored at `index`.
    pub(super) fn new_node_id(index: usize) -> Self {
        let index = index + 1; // index 0 is reserved for the constant 0.
        let index = index << 1; // LSB is used for marking signals as 'inverted'.
        assert!(
            index < (1 << (usize::BITS - 1)),
            "index out of supported range"
        );
        Self { index }
    }

    /// Create a pointer to an input node.
    pub(super) fn new_input_node(index: usize) -> Self {
        let index = index << 1; // LSB is used for marking signals as 'inverted'.
        assert!(
            index < (1 << (usize::BITS - 1)),
            "index out of supported range"
        );
        let input_marker_bit = 1 << (usize::BITS - 1);

        Self {
            index: index | input_marker_bit,
        }
    }

    /// Check if the ID points to an input node.
    pub fn is_input(&self) -> bool {
        let input_marker_bit = 1 << (usize::BITS - 1);
        (self.index & input_marker_bit) != 0
    }

    /// Check if this ID points to the constant 0 or 1.
    pub fn is_constant(&self) -> bool {
        self.is_zero() || self.is_one()
    }

    /// Check if the signal represents the constant 0.
    pub fn is_zero(&self) -> bool {
        self.index == 0
    }

    /// Check if the signal represents the constant 1.
    pub fn is_one(&self) -> bool {
        self.invert().index == 0
    }
}

impl<Node> LogicNetwork<Node>
where
    Node: NetworkNode + Hash + MutNetworkNodeWithReferenceCount + IntoIterator<Item = NodeId>,
{
    /// Insert a node in the storage.
    /// The node should have been normalized beforehand in order to make structural hashing effective. E.g. inputs of commutative nodes should be sorted.
    fn insert_node(&mut self, node: Node) -> NodeId {
        //debug_assert!(
        //    (0..node.num_inputs())
        //        .skip(1)
        //        .all(|i| node.get_input(i - 1) <= node.get_input(i)),
        //    "inputs must be sorted"  // TODO: This should hold only for commutative nodes.
        //);

        let index = self.storage.len();
        self.storage.push(node.clone().into());
        let new_id = NodeId::new_node_id(index);
        let old_node = self.hash_table.insert(node.clone(), new_id);
        debug_assert!(old_node.is_none(), "node already exists");

        // Update reference counters of used nodes.
        // Update references.
        {
            // Iterate over fan-ins.
            node.into_iter()
                .dedup()
                .filter_map(|fan_in_node_id| fan_in_node_id.node_index())
                .for_each(|fan_in_node_index| {
                    self.insert_reverse_edge(new_id, fan_in_node_index);
                })
        }

        new_id
    }

    /// Enables traversing the DAG in opposite direction of the directed edges.
    fn insert_reverse_edge(&mut self, node_id: NodeId, fan_in_node: usize) {
        // Increment the reference counter.
        self.storage[fan_in_node].node.reference();

        // Create an edge from the input node to the new node.
        let list_of_references = &mut self.storage[fan_in_node].references;
        let edge_id = self.references.prepend(list_of_references, node_id);

        let list_of_edges = &mut self.storage[fan_in_node].outgoing_edges;
        self.edges
            .prepend(list_of_edges, (NodeId::new_node_id(fan_in_node), edge_id));
    }

    /// If the node already exists in the graph, replace it with the node ID.
    /// Note: Normalize the node first to increase chances that an already existing equal node can be found.
    fn simplify_node_with_hashtable(&self, node: Node) -> SimplifyResult<Node, NodeId> {
        self.hash_table
            .get(&node)
            .map(|s| SimplifyResult::Simplified(*s, false))
            .unwrap_or(SimplifyResult::new_node(node))
    }
}

impl<Node> ReferenceCounted for LogicNetwork<Node>
where
    Node: NetworkNode<NodeId = NodeId> + NetworkNodeWithReferenceCount,
{
    fn num_references(&self, a: Self::Signal) -> usize {
        if let Some(idx) = a.node_index() {
            self.storage[idx].node.num_references()
        } else {
            // TODO: Reference counts for inputs and constants.
            0
        }
    }
}

impl<Node> LogicNetwork<Node> {
    /// Get the node value of the given `node_id`.
    pub fn get_node(&self, node_id: NodeId) -> GeneralNode<&Node> {
        if let Some(idx) = node_id.node_index() {
            GeneralNode::Node(&self.storage[idx].node)
        } else if let Some(idx) = node_id.input_index() {
            GeneralNode::Input(idx)
        } else {
            assert!(node_id.is_constant());
            GeneralNode::Constant(!node_id.is_zero())
        }
    }

    /// Get a mutable reference to the logic node with the given ID.
    /// Returns `None` if the node is an input, a constant or non-existent.
    pub fn get_logic_node_mut(&mut self, node_id: NodeId) -> Option<&mut Node> {
        node_id.node_index().map(|idx| &mut self.storage[idx].node)
    }
}

impl Edge for NodeId {}

impl EdgeWithInversion for NodeId {
    fn is_inverted(&self) -> bool {
        self.index & 0b1 == 1
    }

    fn invert(self) -> Self {
        Self {
            index: self.index ^ 0b1,
        }
    }

    /// Make the signal non-inverted.
    fn non_inverted(self) -> Self {
        Self {
            index: self.index & !0b1,
        }
    }
}

/// A node, input or constant in the logic network.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub enum GeneralNode<Node> {
    /// A logic node.
    Node(Node),
    /// A primary input.
    Input(usize),
    /// A constant value, either `false` or `true`.
    Constant(bool),
}

impl<Node> GeneralNode<Node> {
    /// Convert to a logic node, if possible.
    pub fn to_logic_node(self) -> Option<Node> {
        match self {
            Self::Node(m) => Some(m),
            _ => None,
        }
    }
    /// Convert to a constant, if possible.
    pub fn to_constant(self) -> Option<bool> {
        match self {
            Self::Constant(c) => Some(c),
            _ => None,
        }
    }
    /// Convert to an index of a primary input.
    pub fn to_input(self) -> Option<usize> {
        match self {
            Self::Input(idx) => Some(idx),
            _ => None,
        }
    }
}

impl<Node> Network for LogicNetwork<Node>
where
    Node: NetworkNode<NodeId = NodeId>,
{
    type Node = Node;
    type NodeId = NodeId;
    type Signal = NodeId;
    type LogicValue = bool;
    type NodeFunction = SmallTruthTable;

    fn num_gates(&self) -> usize {
        self.storage.len()
    }

    fn num_primary_inputs(&self) -> usize {
        self.primary_inputs.len()
    }

    fn num_primary_outputs(&self) -> usize {
        self.primary_outputs.len()
    }

    fn foreach_gate(&self, f: impl Fn(Self::Signal)) {
        (0..self.storage.len()).map(NodeId::new_node_id).for_each(f)
    }

    fn foreach_node(&self, f: impl Fn(&Self::Node)) {
        self.storage.iter().for_each(|n| f(&n.node));
    }

    fn get_constant(&self, value: Self::LogicValue) -> Self::Signal {
        match value {
            false => NodeId::zero(),
            true => NodeId::zero().invert(),
        }
    }

    fn get_constant_value(&self, signal: Self::Signal) -> Option<Self::LogicValue> {
        if signal.is_zero() {
            Some(false)
        } else if signal.is_one() {
            Some(true)
        } else {
            None
        }
    }

    fn get_node_input(&self, node: &Self::Signal, input_index: usize) -> Self::Signal {
        match self.get_node(*node) {
            GeneralNode::Node(n) => n.get_input(input_index),
            GeneralNode::Input(_) => panic!("primary input has no inputs"),
            GeneralNode::Constant(_) => panic!("constant has no inputs"),
        }
    }

    fn get_source_node(&self, signal: &Self::Signal) -> Self::NodeId {
        signal.non_inverted()
    }

    fn get_primary_input(&self, index: usize) -> Self::Signal {
        self.primary_inputs[index]
    }

    fn get_primary_output(&self, index: usize) -> Self::Signal {
        self.primary_outputs[index]
    }

    fn is_constant(&self, signal: Self::Signal) -> bool {
        signal.is_constant()
    }

    fn is_input(&self, signal: Self::Signal) -> bool {
        signal.is_input()
    }

    fn node_function(&self, node: Self::Signal) -> Self::NodeFunction {
        match self.get_node(node) {
            GeneralNode::Node(n) => n.function(),
            GeneralNode::Input(_) => truth_table_library::identity1(),
            GeneralNode::Constant(c) => SmallTruthTable::new(|[]| c),
        }
    }

    fn num_node_inputs(&self, node: &Self::Signal) -> usize {
        match self.get_node(*node) {
            GeneralNode::Node(n) => n.num_inputs(),
            GeneralNode::Input(_) => 0,
            GeneralNode::Constant(_) => 0,
        }
    }

    fn get_node_output(&self, node: &Self::NodeId) -> Self::Signal {
        *node
    }
}

impl<Node> NumInputs for LogicNetwork<Node> {
    fn num_inputs(&self) -> usize {
        self.primary_inputs.len()
    }
}

impl<Node> NumOutputs for LogicNetwork<Node> {
    fn num_outputs(&self) -> usize {
        self.primary_outputs.len()
    }
}

impl<Node> PartialBooleanSystem for LogicNetwork<Node>
where
    Node: NetworkNode<NodeId = NodeId>,
{
    type LiteralId = NodeId;

    type TermId = NodeId;

    fn evaluate_term_partial(&self, term: &Self::TermId, input_values: &[bool]) -> Option<bool> {
        // The logic function of a MIG is always fully specified.
        Some(self.evaluate_term(term, input_values))
    }
}

impl<Node> BooleanSystem for LogicNetwork<Node>
where
    Node: NetworkNode<NodeId = NodeId>,
{
    fn evaluate_term(&self, term: &Self::TermId, input_values: &[bool]) -> bool {
        let simulator = network_simulator::RecursiveSim::new(self);

        assert_eq!(input_values.len(), self.num_primary_inputs());

        let outputs: SmallVec<[_; 1]> = simulator.simulate(&[*term], input_values).collect();

        outputs[0]
    }
}

impl<Node> NetworkEdit for LogicNetwork<Node>
where
    Node: NetworkNode<NodeId = NodeId>
        + Hash
        + Eq
        + MutNetworkNodeWithReferenceCount
        + IntoIterator<Item = NodeId>,
{
    fn create_primary_input(&mut self) -> Self::Signal {
        let idx = self.primary_inputs.len();
        let signal = NodeId::new_input_node(idx);
        self.primary_inputs.push(signal);
        signal
    }

    fn create_primary_output(&mut self, signal: Self::Signal) -> Self::Signal {
        self.primary_outputs.push(signal);
        signal
    }

    fn substitute(&mut self, old: Self::Signal, new: Self::Signal) {
        if let Some(idx) = old.node_index() {
            let list_of_references = &mut self.storage[idx].references;
            let list_of_edges = &mut self.storage[idx].outgoing_edges;
        } else {
            todo!()
        }
    }

    fn create_node(&mut self, node: Self::Node) -> Self::NodeId {
        match self.simplify_node_with_hashtable(node) {
            SimplifyResult::Node(n, inv) => self.insert_node(n).invert_if(inv),
            SimplifyResult::Simplified(node_id, inv) => node_id.invert_if(inv),
        }
    }

    fn foreach_node_mut(&mut self, f: impl Fn(&mut Self::Node)) {
        self.storage.iter_mut().for_each(|n| f(&mut n.node))
    }
}

// impl<Node> GraphBase for LogicNetwork<Node> {
//     type EdgeId = NodeId;
//     type NodeId = NodeId;
// }
//
// impl<Node> IntoNeighbors for &LogicNetwork<Node> {
//     type Neighbors;
//
//     fn neighbors(self, a: Self::NodeId) -> Self::Neighbors {
//         todo!()
//     }
// }
//
// impl<Node> IntoNeighborsDirected for &LogicNetwork<Node> {
//     type NeighborsDirected;
//
//     fn neighbors_directed(
//         self,
//         n: Self::NodeId,
//         d: petgraph::Direction,
//     ) -> Self::NeighborsDirected {
//         todo!()
//     }
// }
//
// impl<Node> IntoNodeIdentifiers for &LogicNetwork<Node> {
//     type NodeIdentifiers;
//
//     fn node_identifiers(self) -> Self::NodeIdentifiers {
//         todo!()
//     }
// }
//