// Panopticon - A libre program analysis library for machine code
// Copyright (C) 2014-2018  The Panopticon Developers
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

//! Generic table for mapping complex types to integers.
//!
//! This module implements `Table<K, V>`. It is used in `Function` for performace optimization.

use std::collections::HashMap;
use std::hash::Hash;
use std::clone::Clone;
use std::fmt::Debug;

use quickcheck::{Arbitrary,Gen};

use {Result,Str,Variable};

#[macro_export]
macro_rules! define_table_ref {
    ( $name:ident ) => {
        /// Table index
        #[derive(Hash,Clone,Copy,Debug,PartialEq,Eq,PartialOrd,Ord)]
        pub struct $name (usize);

        impl $name {
            /// Returns the numberic index.
            pub fn index(&self) -> usize { self.0 }

            /// Creates a new index.
            pub fn new(i: usize) -> $name { $name(i) }
        }

        impl From<usize> for $name {
            fn from(i: usize) -> $name { $name::new(i) }
        }

        impl Into<usize> for $name {
            fn into(self) -> usize { self.index() }
        }
    };
}

define_table_ref!(NameRef);
define_table_ref!(SegmentRef);
define_table_ref!(StrRef);

/// A SSA variable name.
#[derive(Hash,Clone,Debug,PartialEq,Eq,PartialOrd,Ord)]
pub struct Name {
    base: Str,
    subscript: Option<usize>,
}

impl Name {
    /// Create a new SSA variable name.
    pub fn new<S: Into<Option<usize>> + Sized>(s: Str, o: S) -> Name {
        Name{
            base: s,
            subscript: o.into(),
        }
    }

    /// Returns the SSA variable name without subscript.
    pub fn base<'a>(&'a self) -> &'a Str { &self.base }

    /// Returns the numeric subscript of this SSA name.
    pub fn subscript(&self) -> Option<usize> { self.subscript }
}

/// Table mapping hashable values to and from numeric indices.
#[derive(Debug,Clone)]
pub struct Table<V: Hash + Debug + Clone + Eq,R: Clone + Copy + From<usize> + Into<usize>> {
    refs: HashMap<V,R>,
    values: Vec<V>,
    facade: bool
}

impl<V: Hash + Debug + Clone + Eq,R: Clone + Copy + From<usize> + Into<usize>> Table<V,R> {
    /// Creates a new table with initial capacity `cap`.
    pub fn with_capacity(cap: usize) -> Table<V,R> {
        Table{
            refs: HashMap::with_capacity(cap),
            values: Vec::with_capacity(cap),
            facade: false,
        }
    }

    /// Inserts `v` into the table, returing its index. If `v` already exists in the table its not inserted again.
    pub fn insert(&mut self, v: &V) -> R {
        if let Some(i) = self.refs.get(v).cloned() {
            i
        } else {
            let i = self.values.len();
            self.refs.insert(v.clone(),i.into());
            self.values.push(v.clone());
            i.into()
        }
    }

    /// Returns the index of `s` in the table.
    pub fn index(&self, s: &V) -> Result<R> {
        match self.refs.get(s).map(|&x| x) {
            Some(s) => Ok(s),
            None if self.facade => Ok(0.into()),
            None => Err("name not part of this set".into()),
        }
    }

    /// Returns the value of the entry with index `i`.
    pub fn value<'a>(&'a self, i: R) -> Result<&'a V> {
        if self.facade {
            Ok(&self.values[0])
        } else {
            Ok(&self.values[i.into()])
        }
    }

    /// Number of entries in the table.
    pub fn len(&self) -> usize {
        self.values.len()
    }

    /// Creates a dummy tables that maps every index to `v`.
    #[cfg(test)]
    pub fn facade(v: &V) -> Table<V,R> {
        let mut ret = Table{
            refs: HashMap::default(),
            values: Vec::default(),
            facade: true,
        };

        ret.insert(v);
        ret
    }
}

impl Table<Name,NameRef> {
    /// Shortcut for getting the index of the `Name` with the same base as `idx` but subscript
    /// `None`.
    pub fn base_name(&self, idx: NameRef) -> Result<NameRef> {
        let name = self.value(idx)?;
        match name {
            &Name{ ref base, subscript: Some(_) } => {
                self.index(&Name{ base: base.clone(), subscript: None })
            }
            &Name{ subscript: None,.. } => {
                Ok(idx)
            }
        }
    }

    /// Create a new variable with `name` and `subscript`, inserting the name if it's not present
    /// in the table.
    pub fn var<S: Into<Str>>(&mut self, name: S, subscript: Option<usize>, bits: usize) -> Result<Variable> {
        let i = self.insert(&Name::new(name.into(),subscript));
        Variable::new(i,bits)
    }
}

impl Table<Name,SegmentRef> {
    /// Shortcut for getting the index of the `Segment` with the same base as `idx` but subscript
    /// `None`.
    pub fn base_name(&self, idx: SegmentRef) -> Result<SegmentRef> {
        let seg = self.value(idx)?;
        match seg {
            &Name{ ref base, subscript: Some(_) } => {
                self.index(&Name{ base: base.clone(), subscript: None })
            }
            &Name{ subscript: None,.. } => {
                Ok(idx)
            }
        }
    }
}

impl<V: Hash + Debug + Clone + Eq,R: Clone + Copy + From<usize> + Into<usize>> Default for Table<V,R> {
    fn default() -> Table<V,R> {
        Table{
            refs: HashMap::default(),
            values: Vec::default(),
            facade: false,
        }
    }
}

impl Arbitrary for NameRef {
    fn arbitrary<G: Gen>(g: &mut G) -> Self {
        NameRef::new(g.gen_range(0,100))
    }
}

impl Arbitrary for StrRef {
    fn arbitrary<G: Gen>(g: &mut G) -> Self {
        StrRef::new(g.gen_range(0,100))
    }
}

/// String table.
pub type Strings = Table<Str,StrRef>;

/// Name table for RREIL code.
pub type Names = Table<Name,NameRef>;

/// Segment table RREIL code.
pub type Segments = Table<Name,SegmentRef>;
