// 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

//! IL Statement encoding.
//!
//! To save space IL statements inside a function can be "compressed" by enconding the into a more
//! space conserving, binary representation.
//!
//! This module handles convertion between those two.

use std::ops::Range;
use std::slice::Iter;
use std::usize;
use std::borrow::Cow;

use {
    Mnemonic,BasicBlock,Names,Strings,
    Result,
    Statement,Bitcode,BitcodeIter,Segments
};

/// Iterator over IL statements
pub enum StatementsIter<'a> {
    Bitcode(BitcodeIter<'a>),
    Vector(Iter<'a,Statement>),
}

impl<'a> Iterator for StatementsIter<'a> {
    type Item = Cow<'a,Statement>;

    fn next(&mut self) -> Option<Cow<'a,Statement>> {
        match self {
            &mut StatementsIter::Bitcode(ref mut bt) => bt.next().map(|x| Cow::Owned(x)),
            &mut StatementsIter::Vector(ref mut v) => v.next().map(|x| Cow::Borrowed(x)),
        }
    }
}

/// Signaling for `Statement::rewrite`.
pub enum RewriteControl {
    /// Continue with the next IL statement.
    Continue,
    /// Stop iteration.
    Break,
}

/// A possibly compressed IL sequence.
#[derive(Clone, Debug)]
pub enum Statements {
    /// Compressed binary encoding for IL statements.
    Bitcode(Bitcode),
    /// Uncompressed statement vector for processing.
    Vector(Vec<Statement>),
}

impl Default for Statements {
    fn default() -> Statements {
        Statements::Vector(Vec::new())
    }
}

impl Statements {
    /// Creates an iterator over IL statements inside `rgn`.
    pub fn iter_range<'a>(&'a self, rgn: Range<usize>) -> StatementsIter<'a> {
        match self {
            &Statements::Bitcode(ref bt) => StatementsIter::Bitcode(bt.iter_range(rgn)),
            &Statements::Vector(ref v) => StatementsIter::Vector(v[rgn].iter()),
        }
    }

    /// Mutate all IL statements in `rgn` using `f`.
    pub fn rewrite<F: FnMut(&mut Statement, &mut Names, &mut Strings, &mut Segments) -> Result<RewriteControl> + Sized>(&mut self, rgn: Range<usize>, names: &mut Names, strings: &mut Strings, segments: &mut Segments, mut f: F) -> Result<Range<usize>> {
        match self {
            &mut Statements::Bitcode(ref mut bt) => bt.rewrite(rgn,names,strings,segments,f),
            &mut Statements::Vector(ref mut v) => {
                for stmt in v[rgn.clone()].iter_mut() {
                    match f(stmt,names,strings,segments)? {
                        RewriteControl::Break => { break; }
                        RewriteControl::Continue => { continue; }
                    }
                }
                Ok(rgn)
            }
        }
    }

    /// Inserts IL statements `stmts` at position `pos`.
    pub fn insert(&mut self, pos: usize, stmts: Vec<Statement>) -> Result<Range<usize>> {
        match self {
            &mut Statements::Bitcode(ref mut bt) => bt.insert(pos,stmts),
            &mut Statements::Vector(ref mut v) => {
                let rgn = pos..(pos + stmts.len());
                for (i,stmt) in stmts.into_iter().enumerate() {
                    v.insert(i + pos,stmt);
                }
                Ok(rgn)
            }
        }
    }

    /// Removes IL statements in `rgn`
    pub fn remove(&mut self, rgn: Range<usize>) {
        match self {
            &mut Statements::Bitcode(ref mut bt) => bt.remove(rgn),
            &mut Statements::Vector(ref mut v) => { v.drain(rgn); }
        }
    }

    /// Appends all IL statements in `i` to the instance.
    pub fn append<I: IntoIterator<Item=Statement> + Sized>(&mut self, i: I) -> Result<Range<usize>> {
        match self {
            &mut Statements::Bitcode(ref mut bt) => bt.append(i),
            &mut Statements::Vector(ref mut v) => {
                let fst = v.len();
                v.extend(i);
                Ok(fst..v.len())
            }
        }
    }

    /// Upper end of the allowable IL statement range.
    pub fn len(&self) -> usize {
        match self {
            &Statements::Bitcode(ref bt) => bt.num_bytes(),
            &Statements::Vector(ref v) => v.len(),
        }
    }

    /// Compress IL statements into bitcode, updates `basic_blocks` and `mnemonics` statement
    /// ranges.
    pub fn pack(&mut self, basic_blocks: &mut Vec<BasicBlock>, mnemonics: &mut Vec<Mnemonic>) -> Result<()> {
        *self = match self {
            &mut Statements::Bitcode(_) => { return Ok(()); },
            &mut Statements::Vector(ref mut vec) => {
                let mut bc = Bitcode::with_capacity(vec.len() * 15);

                for bb in basic_blocks.iter_mut() {
                    let mut old_idx = bb.statements.start;
                    let mut new_idx = bc.num_bytes();
                    let new_start = new_idx;
                    let mne_rgn = bb.mnemonics.start.index()..bb.mnemonics.end.index();

                    for mne in mnemonics[mne_rgn].iter_mut() {
                        let old_rgn = old_idx..(old_idx + mne.statement_count);
                        let new_rgn = bc.append(vec[old_rgn].iter().cloned())?;

                        old_idx += mne.statement_count;
                        mne.statement_count = new_rgn.end - new_rgn.start;
                        new_idx += mne.statement_count;
                    }

                    assert_eq!(bb.statements.end, old_idx);
                    bb.statements = new_start..new_idx;
                }

                Statements::Bitcode(bc)
            }
        };

        Ok(())
    }

    /// Uncompress IL statements into bitcode, updates `basic_blocks` and `mnemonics` statement
    /// ranges.
     pub fn unpack(&mut self, basic_blocks: &mut Vec<BasicBlock>, mnemonics: &mut Vec<Mnemonic>) -> Result<()> {
        *self = match self {
            &mut Statements::Vector(_) => { return Ok(()); }
            &mut Statements::Bitcode(ref bc) => {
                let mut vec = Vec::with_capacity(bc.num_bytes() / 15);

                for bb in basic_blocks.iter_mut() {
                    let mut old_idx = bb.statements.start;
                    let mut new_idx = vec.len();
                    let new_start = new_idx;
                    let mne_rgn = bb.mnemonics.start.index()..bb.mnemonics.end.index();

                    for mne in mnemonics[mne_rgn].iter_mut() {
                        let old_rgn = old_idx..(old_idx + mne.statement_count);

                        vec.extend(bc.iter_range(old_rgn));
                        old_idx += mne.statement_count;
                        mne.statement_count = vec.len() - new_idx;
                        new_idx += mne.statement_count;
                    }

                    assert_eq!(bb.statements.end, old_idx);
                    bb.statements = new_start..new_idx;
                }

                Statements::Vector(vec)
            }
        };

        Ok(())
    }
}