<commit_before>#![cfg(feature = "use_std")]

use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Iterator;

/// Creates a new `GroupingMap` from `iter`
pub fn new<I, K, V>(iter: I) -> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    GroupingMap { iter }
}

/// `GroupingMap` is an intermediate struct for efficient group-and-fold operations.
/// It groups elements by their key and at the same time fold each group
/// using some aggregating operation.
///
/// No method on this struct performs temporary allocations.
pub struct GroupingMap<I> {
    iter: I,
}

impl<I, K, V> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    /// This is the generic way to perform any operations on a `GroupingMap`.
    /// It's suggested to use this method only to implement custom operations
    /// when the already provided ones are not enough.
    ///
    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in an `HashMap`.
    ///
    /// The `operation` function is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group if there is currently one;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being aggregated;
    /// If `operation` returns `Some(element)` then the accumulator is updated with `element`,
    /// otherwise the previous accumulation is discarded.
    ///
    /// Return a `HashMap` associating the key of each group with the result of aggregation of the group elements.
    /// If there's no result then there won't be an entry associated to that key.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let data = vec![2, 8, 5, 7, 9, 0, 4, 10];
    /// let lookup = data.into_iter()
    ///     .map(|n| (n % 4, n))
    ///     .into_grouping_map()
    ///     .aggregate(|acc, _, val| {
    ///         if val == 0 || val == 10 {
    ///             None
    ///         } else {
    ///             Some(acc.unwrap_or(0) + val)
    ///         }
    ///     });
    ///
    /// assert_eq!(lookup[&0], 4);        // 0 resets the accumulator so only 4 is summed
    /// assert_eq!(lookup[&1], 14);       // 5 + 9
    /// assert_eq!(lookup.get(&2), None); // 10 resets the accumulator and nothing is summed afterward
    /// assert_eq!(lookup[&3], 7);        // 7
    /// assert_eq!(lookup.len(), 3);      // The final keys are only 0, 1 and 2
    /// ```
    pub fn aggregate<FO, R>(self, mut operation: FO) -> HashMap<K, R>
    where
        FO: FnMut(Option<R>, &K, V) -> Option<R>,
    {
        let mut destination_map = HashMap::new();

        for (key, val) in self.iter {
            let acc = destination_map.remove(&key);
            if let Some(op_res) = operation(acc, &key, val) {
                destination_map.insert(key, op_res);
            }
        }

        destination_map
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// `init` is the value from which will be cloned the initial value of each accumulator.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold(0, |acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);
    /// assert_eq!(lookup[&1], 12);
    /// assert_eq!(lookup[&2], 7);
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold<FO, R>(self, init: R, mut operation: FO) -> HashMap<K, R>
    where
        R: Clone,
        FO: FnMut(R, &K, V) -> R,
    {
        self.aggregate(|acc, key, val| {
            let acc = acc.unwrap_or_else(|| init.clone());
            Some(operation(acc, key, val))
        })
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// This is similar to [`fold`] but the initial value of the accumulator is the first element of the group.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// [`fold`]: #tymethod.fold
    pub fn fold_first<FO>(self, mut operation: FO) -> HashMap<K, V>
    where
        FO: FnMut(V, &K, V) -> V,
    {
        self.aggregate(|acc, key, val| {
            Some(match acc {
                Some(acc) => operation(acc, key, val),
                None => val,
            })
        })
    }

    /// Groups elements from the `GroupingMap` source by key and collects the elements of each group in
    /// an instance of `C`. The iteration order is preserved when inserting elements.
    ///
    /// Return a `HashMap` associating the key of each group with the collection containing the elements of that group.
    pub fn collect<C>(self) -> HashMap<K, C>
    where
        C: Default + Extend<V>,
    {
        self.aggregate(|acc, _, v| {
            let mut acc = acc.unwrap_or_else(C::default);
            acc.extend(Some(v));
            Some(acc)
        })
    }

    /// Groups elements from the `GroupingMap` source by key and counts them.
    ///
    /// Return a `HashMap` associating the key of each group with the number of elements in that group.
    pub fn count(self) -> HashMap<K, usize> {
        self.fold(0, |acc, _, _| acc + 1)
    }
}
<commit_msg>Improved doctest for GroupingMap::fold<commit_after>#![cfg(feature = "use_std")]

use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Iterator;

/// Creates a new `GroupingMap` from `iter`
pub fn new<I, K, V>(iter: I) -> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    GroupingMap { iter }
}

/// `GroupingMap` is an intermediate struct for efficient group-and-fold operations.
/// It groups elements by their key and at the same time fold each group
/// using some aggregating operation.
///
/// No method on this struct performs temporary allocations.
pub struct GroupingMap<I> {
    iter: I,
}

impl<I, K, V> GroupingMap<I>
where
    I: Iterator<Item = (K, V)>,
    K: Hash + Eq,
{
    /// This is the generic way to perform any operations on a `GroupingMap`.
    /// It's suggested to use this method only to implement custom operations
    /// when the already provided ones are not enough.
    ///
    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in an `HashMap`.
    ///
    /// The `operation` function is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group if there is currently one;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being aggregated;
    /// If `operation` returns `Some(element)` then the accumulator is updated with `element`,
    /// otherwise the previous accumulation is discarded.
    ///
    /// Return a `HashMap` associating the key of each group with the result of aggregation of the group elements.
    /// If there's no result then there won't be an entry associated to that key.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let data = vec![2, 8, 5, 7, 9, 0, 4, 10];
    /// let lookup = data.into_iter()
    ///     .map(|n| (n % 4, n))
    ///     .into_grouping_map()
    ///     .aggregate(|acc, _, val| {
    ///         if val == 0 || val == 10 {
    ///             None
    ///         } else {
    ///             Some(acc.unwrap_or(0) + val)
    ///         }
    ///     });
    ///
    /// assert_eq!(lookup[&0], 4);        // 0 resets the accumulator so only 4 is summed
    /// assert_eq!(lookup[&1], 14);       // 5 + 9
    /// assert_eq!(lookup.get(&2), None); // 10 resets the accumulator and nothing is summed afterward
    /// assert_eq!(lookup[&3], 7);        // 7
    /// assert_eq!(lookup.len(), 3);      // The final keys are only 0, 1 and 2
    /// ```
    pub fn aggregate<FO, R>(self, mut operation: FO) -> HashMap<K, R>
    where
        FO: FnMut(Option<R>, &K, V) -> Option<R>,
    {
        let mut destination_map = HashMap::new();

        for (key, val) in self.iter {
            let acc = destination_map.remove(&key);
            if let Some(op_res) = operation(acc, &key, val) {
                destination_map.insert(key, op_res);
            }
        }

        destination_map
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// `init` is the value from which will be cloned the initial value of each accumulator.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// ```
    /// use itertools::Itertools;
    ///
    /// let lookup = (1..=7)
    ///     .map(|n| (n % 3, n))
    ///     .into_grouping_map()
    ///     .fold(0, |acc, _, val| acc + val);
    ///
    /// assert_eq!(lookup[&0], 9);   // 3 + 6
    /// assert_eq!(lookup[&1], 12);  // 1 + 4 + 7
    /// assert_eq!(lookup[&2], 7);   // 2 + 5
    /// assert_eq!(lookup.len(), 3);
    /// ```
    pub fn fold<FO, R>(self, init: R, mut operation: FO) -> HashMap<K, R>
    where
        R: Clone,
        FO: FnMut(R, &K, V) -> R,
    {
        self.aggregate(|acc, key, val| {
            let acc = acc.unwrap_or_else(|| init.clone());
            Some(operation(acc, key, val))
        })
    }

    /// Groups elements from the `GroupingMap` source by key and applies `operation` to the elements
    /// of each group sequentially, passing the previously accumulated value, a reference to the key
    /// and the current element as arguments, and stores the results in a new map.
    ///
    /// This is similar to [`fold`] but the initial value of the accumulator is the first element of the group.
    ///
    /// `operation` is a function that is invoked on each element with the following parameters:
    ///  - the current value of the accumulator of the group;
    ///  - a reference to the key of the group this element belongs to;
    ///  - the element from the source being accumulated.
    ///
    /// Return a `HashMap` associating the key of each group with the result of folding the group elements.
    ///
    /// [`fold`]: #tymethod.fold
    pub fn fold_first<FO>(self, mut operation: FO) -> HashMap<K, V>
    where
        FO: FnMut(V, &K, V) -> V,
    {
        self.aggregate(|acc, key, val| {
            Some(match acc {
                Some(acc) => operation(acc, key, val),
                None => val,
            })
        })
    }

    /// Groups elements from the `GroupingMap` source by key and collects the elements of each group in
    /// an instance of `C`. The iteration order is preserved when inserting elements.
    ///
    /// Return a `HashMap` associating the key of each group with the collection containing the elements of that group.
    pub fn collect<C>(self) -> HashMap<K, C>
    where
        C: Default + Extend<V>,
    {
        self.aggregate(|acc, _, v| {
            let mut acc = acc.unwrap_or_else(C::default);
            acc.extend(Some(v));
            Some(acc)
        })
    }

    /// Groups elements from the `GroupingMap` source by key and counts them.
    ///
    /// Return a `HashMap` associating the key of each group with the number of elements in that group.
    pub fn count(self) -> HashMap<K, usize> {
        self.fold(0, |acc, _, _| acc + 1)
    }
}
<|endoftext|>


<commit_before>use dynasmrt::x64::Assembler;
use dynasmrt::{DynasmApi, DynasmLabelApi};

use super::*;
use super::analysis::{BoundsAnalysis, AbstractInterpreter, NoAnalysis};
use common::Count;
use peephole;

/// Program forms that can be JIT compiled.
pub trait JitCompilable {
    /// Compile the given program into the peephole AST to prepare for JIT compilation.
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R;

    /// JIT compile the given program.
    fn jit_compile(&self, checked: bool) -> Program {
        self.with_peephole(|ast| compile(ast, checked))
    }
}

dynasm!(asm
    ; .alias pointer, r12
    ; .alias mem_start, r13
    ; .alias mem_limit, r14
    ; .alias rts, r15
);

/// Compiles peephole-optimized AST to x64 machine code.
///
/// Uses the `dynasmrt` assembler
pub fn compile(program: &peephole::Program, checked: bool) -> Program {
    if checked {
        let mut compiler = Compiler::<AbstractInterpreter>::new(program, true);
        compiler.compile(program);
        compiler.into_program()
    } else {
        let mut compiler = Compiler::<NoAnalysis>::new(program, false);
        compiler.compile(program);
        compiler.into_program()
    }
}

/// The compiler state.
struct Compiler<B: BoundsAnalysis> {
    /// The underlying assembler.
    asm: Assembler,
    /// The offset of the starting instruction for the object function.
    start: dynasmrt::AssemblyOffset,
    /// Whether we are emitting bounds checks.
    checked: bool,
    /// Abstract interpreter for bounds checking analysis.
    interpreter: B,
}

impl<B: BoundsAnalysis> Compiler<B> {
    fn new(program: &peephole::Program, checked: bool) -> Self {
        let asm = Assembler::new();
        let start = asm.offset();

        let mut result = Compiler {
            asm: asm,
            start: start,
            checked: checked,
            interpreter: B::new(program),
        };

        result.emit_prologue();

        result
    }

    fn into_program(mut self) -> Program {
        self.emit_epilogue();

        Program {
            code: self.asm.finalize().unwrap(),
            start: self.start,
        }
    }

    fn emit_prologue(&mut self) {
        dynasm!(self.asm
            ; push r12
            ; push r13
            ; push r14
            ; push r15
            ; mov pointer, rcx      // first argument
            ; mov mem_start, rcx
            ; mov mem_limit, rcx
            ; add mem_limit, rdx    // second argument
            ; mov rts, r8           // third argument
        );
    }

    fn emit_epilogue(&mut self) {
        dynasm!(self.asm
            ; mov rax, rts::OKAY as i32
            ; jmp ->finish

            ; ->underflow:
            ; mov rax, rts::UNDERFLOW as i32
            ; jmp ->finish

            ; ->overflow:
            ; mov rax, rts::OVERFLOW as i32

            ; ->finish:
            ; pop r15
            ; pop r14
            ; pop r13
            ; pop r12
            ; ret
        );
    }

    fn compile(&mut self, program: &[peephole::Statement]) {
        for stm in program {
            self.compile_statement(stm);
        }
    }

    fn compile_statement(&mut self, stm: &peephole::Statement) {
        use peephole::Statement::*;
        use common::Instruction::*;

        match *stm {
            Instr(Right(count)) => {
                let proved = self.interpreter.move_right(count);

                dynasm!(self.asm
                    ;; self.load_pos_offset(count, proved)
                    ; add pointer, rax
                );
            }

            Instr(Left(count)) => {
                let proved = self.interpreter.move_left(count);

                dynasm!(self.asm
                    ;; self.load_neg_offset(count, proved)
                    ; sub pointer, rax
                );
            }

            Instr(Add(count)) => {
                dynasm!(self.asm
                    ; add [pointer], BYTE count as i8
                );
            }

            Instr(In) => {
                dynasm!(self.asm
                    ;; self.rts_call(rts::RtsState::read as _)
                    ; mov [pointer], al
                );
            }

            Instr(Out) => {
                dynasm!(self.asm
                    ; xor rdx, rdx
                    ; mov dl, [pointer]
                    ;; self.rts_call(rts::RtsState::write as _)
                );
            }

            Instr(SetZero) => {
                dynasm!(self.asm
                    ; mov BYTE [pointer], 0
                )
            }

            Instr(FindZeroRight(skip)) => {
                self.interpreter.reset_right();

                dynasm!(self.asm
                    ; jmp >end_loop
                    ; begin_loop:
                    ;; self.load_pos_offset(skip, false)
                    ; add pointer, rax
                    ; end_loop:
                    ; cmp BYTE [pointer], 0
                    ; jnz <begin_loop
                )
            }

            Instr(FindZeroLeft(skip)) => {
                self.interpreter.reset_left();

                dynasm!(self.asm
                    ; jmp >end_loop
                    ; begin_loop:
                    ;; self.load_neg_offset(skip, false)
                    ; sub pointer, rax
                    ; end_loop:
                    ; cmp BYTE [pointer], 0
                    ; jnz <begin_loop
                )
            }

            Instr(OffsetAddRight(offset)) => {
                let proved = self.interpreter.check_right(offset);

                dynasm!(self.asm
                    ; cmp BYTE [pointer], 0
                    ; jz >skip
                    ;; self.load_pos_offset(offset, proved)
                    ; mov cl, BYTE [pointer]
                    ; mov BYTE [pointer], 0
                    ; add BYTE [pointer + rax], cl
                    ; skip:
                );
            }

            Instr(OffsetAddLeft(offset)) => {
                let proved = self.interpreter.check_left(offset);

                dynasm!(self.asm
                    ; cmp BYTE [pointer], 0
                    ; jz >skip
                    ;; self.load_neg_offset(offset, proved)
                    ; mov cl, BYTE [pointer]
                    ; mov BYTE [pointer], 0
                    ; neg rax
                    ; add BYTE [pointer + rax], cl
                    ; skip:
                );
            }

            Instr(JumpZero(_)) | Instr(JumpNotZero(_)) =>
                panic!("unexpected jump instruction"),

            Loop(ref body) => {
                let begin_label = self.asm.new_dynamic_label();
                let end_label   = self.asm.new_dynamic_label();

                self.interpreter.enter_loop(body);

                dynasm!(self.asm
                    ; jmp =>end_label
                    ; =>begin_label
                    ;; self.compile(body)
                    ; =>end_label
                    ; cmp BYTE [pointer], 0
                    ; jnz =>begin_label
                );

                self.interpreter.leave_loop();
            }
        }
    }

    fn rts_call(&mut self, fun: i64) {
        dynasm!(self.asm
            ; mov rax, QWORD fun
            ; mov rcx, rts
            ; sub rsp, BYTE 0x28
            ; call rax
            ; add rsp, BYTE 0x28
        );
    }

    #[inline]
    fn load_constant(&mut self, count: Count) {
        if count as i32 as Count == count {
            dynasm!(self.asm
                ; mov rax, DWORD count as i32
            );
        } else {
            dynasm!(self.asm
                ; mov rax, QWORD count as i64
            );
        }
    }

    #[inline]
    fn load_pos_offset(&mut self, offset: Count, proved: bool) {
        self.load_constant(offset);

        if self.checked && !proved {
            dynasm!(self.asm
                ; mov rcx, mem_limit
                ; sub rcx, pointer
                ; cmp rcx, rax
                ; jle ->overflow
            );
        }
    }

    #[inline]
    fn load_neg_offset(&mut self, offset: Count, proved: bool) {
        self.load_constant(offset);

        if self.checked && !proved {
            dynasm!(self.asm
                ; mov rcx, pointer
                ; sub rcx, mem_start
                ; cmp rcx, rax
                ; jl ->underflow
            );
        }
    }
}

impl JitCompilable for Box<peephole::Program> {
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R
    {
        k(self)
    }
}

impl<T: peephole::PeepholeCompilable + ?Sized> JitCompilable for T {
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R
    {
        k(&self.peephole_compile())
    }
}
<commit_msg>Changed JitCompilable impl to not involve Box.<commit_after>use dynasmrt::x64::Assembler;
use dynasmrt::{DynasmApi, DynasmLabelApi};

use super::*;
use super::analysis::{BoundsAnalysis, AbstractInterpreter, NoAnalysis};
use common::Count;
use peephole;

/// Program forms that can be JIT compiled.
pub trait JitCompilable {
    /// Compile the given program into the peephole AST to prepare for JIT compilation.
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R;

    /// JIT compile the given program.
    fn jit_compile(&self, checked: bool) -> Program {
        self.with_peephole(|ast| compile(ast, checked))
    }
}

dynasm!(asm
    ; .alias pointer, r12
    ; .alias mem_start, r13
    ; .alias mem_limit, r14
    ; .alias rts, r15
);

/// Compiles peephole-optimized AST to x64 machine code.
///
/// Uses the `dynasmrt` assembler
pub fn compile(program: &peephole::Program, checked: bool) -> Program {
    if checked {
        let mut compiler = Compiler::<AbstractInterpreter>::new(program, true);
        compiler.compile(program);
        compiler.into_program()
    } else {
        let mut compiler = Compiler::<NoAnalysis>::new(program, false);
        compiler.compile(program);
        compiler.into_program()
    }
}

/// The compiler state.
struct Compiler<B: BoundsAnalysis> {
    /// The underlying assembler.
    asm: Assembler,
    /// The offset of the starting instruction for the object function.
    start: dynasmrt::AssemblyOffset,
    /// Whether we are emitting bounds checks.
    checked: bool,
    /// Abstract interpreter for bounds checking analysis.
    interpreter: B,
}

impl<B: BoundsAnalysis> Compiler<B> {
    fn new(program: &peephole::Program, checked: bool) -> Self {
        let asm = Assembler::new();
        let start = asm.offset();

        let mut result = Compiler {
            asm: asm,
            start: start,
            checked: checked,
            interpreter: B::new(program),
        };

        result.emit_prologue();

        result
    }

    fn into_program(mut self) -> Program {
        self.emit_epilogue();

        Program {
            code: self.asm.finalize().unwrap(),
            start: self.start,
        }
    }

    fn emit_prologue(&mut self) {
        dynasm!(self.asm
            ; push r12
            ; push r13
            ; push r14
            ; push r15
            ; mov pointer, rcx      // first argument
            ; mov mem_start, rcx
            ; mov mem_limit, rcx
            ; add mem_limit, rdx    // second argument
            ; mov rts, r8           // third argument
        );
    }

    fn emit_epilogue(&mut self) {
        dynasm!(self.asm
            ; mov rax, rts::OKAY as i32
            ; jmp ->finish

            ; ->underflow:
            ; mov rax, rts::UNDERFLOW as i32
            ; jmp ->finish

            ; ->overflow:
            ; mov rax, rts::OVERFLOW as i32

            ; ->finish:
            ; pop r15
            ; pop r14
            ; pop r13
            ; pop r12
            ; ret
        );
    }

    fn compile(&mut self, program: &[peephole::Statement]) {
        for stm in program {
            self.compile_statement(stm);
        }
    }

    fn compile_statement(&mut self, stm: &peephole::Statement) {
        use peephole::Statement::*;
        use common::Instruction::*;

        match *stm {
            Instr(Right(count)) => {
                let proved = self.interpreter.move_right(count);

                dynasm!(self.asm
                    ;; self.load_pos_offset(count, proved)
                    ; add pointer, rax
                );
            }

            Instr(Left(count)) => {
                let proved = self.interpreter.move_left(count);

                dynasm!(self.asm
                    ;; self.load_neg_offset(count, proved)
                    ; sub pointer, rax
                );
            }

            Instr(Add(count)) => {
                dynasm!(self.asm
                    ; add [pointer], BYTE count as i8
                );
            }

            Instr(In) => {
                dynasm!(self.asm
                    ;; self.rts_call(rts::RtsState::read as _)
                    ; mov [pointer], al
                );
            }

            Instr(Out) => {
                dynasm!(self.asm
                    ; xor rdx, rdx
                    ; mov dl, [pointer]
                    ;; self.rts_call(rts::RtsState::write as _)
                );
            }

            Instr(SetZero) => {
                dynasm!(self.asm
                    ; mov BYTE [pointer], 0
                )
            }

            Instr(FindZeroRight(skip)) => {
                self.interpreter.reset_right();

                dynasm!(self.asm
                    ; jmp >end_loop
                    ; begin_loop:
                    ;; self.load_pos_offset(skip, false)
                    ; add pointer, rax
                    ; end_loop:
                    ; cmp BYTE [pointer], 0
                    ; jnz <begin_loop
                )
            }

            Instr(FindZeroLeft(skip)) => {
                self.interpreter.reset_left();

                dynasm!(self.asm
                    ; jmp >end_loop
                    ; begin_loop:
                    ;; self.load_neg_offset(skip, false)
                    ; sub pointer, rax
                    ; end_loop:
                    ; cmp BYTE [pointer], 0
                    ; jnz <begin_loop
                )
            }

            Instr(OffsetAddRight(offset)) => {
                let proved = self.interpreter.check_right(offset);

                dynasm!(self.asm
                    ; cmp BYTE [pointer], 0
                    ; jz >skip
                    ;; self.load_pos_offset(offset, proved)
                    ; mov cl, BYTE [pointer]
                    ; mov BYTE [pointer], 0
                    ; add BYTE [pointer + rax], cl
                    ; skip:
                );
            }

            Instr(OffsetAddLeft(offset)) => {
                let proved = self.interpreter.check_left(offset);

                dynasm!(self.asm
                    ; cmp BYTE [pointer], 0
                    ; jz >skip
                    ;; self.load_neg_offset(offset, proved)
                    ; mov cl, BYTE [pointer]
                    ; mov BYTE [pointer], 0
                    ; neg rax
                    ; add BYTE [pointer + rax], cl
                    ; skip:
                );
            }

            Instr(JumpZero(_)) | Instr(JumpNotZero(_)) =>
                panic!("unexpected jump instruction"),

            Loop(ref body) => {
                let begin_label = self.asm.new_dynamic_label();
                let end_label   = self.asm.new_dynamic_label();

                self.interpreter.enter_loop(body);

                dynasm!(self.asm
                    ; jmp =>end_label
                    ; =>begin_label
                    ;; self.compile(body)
                    ; =>end_label
                    ; cmp BYTE [pointer], 0
                    ; jnz =>begin_label
                );

                self.interpreter.leave_loop();
            }
        }
    }

    fn rts_call(&mut self, fun: i64) {
        dynasm!(self.asm
            ; mov rax, QWORD fun
            ; mov rcx, rts
            ; sub rsp, BYTE 0x28
            ; call rax
            ; add rsp, BYTE 0x28
        );
    }

    #[inline]
    fn load_constant(&mut self, count: Count) {
        if count as i32 as Count == count {
            dynasm!(self.asm
                ; mov rax, DWORD count as i32
            );
        } else {
            dynasm!(self.asm
                ; mov rax, QWORD count as i64
            );
        }
    }

    #[inline]
    fn load_pos_offset(&mut self, offset: Count, proved: bool) {
        self.load_constant(offset);

        if self.checked && !proved {
            dynasm!(self.asm
                ; mov rcx, mem_limit
                ; sub rcx, pointer
                ; cmp rcx, rax
                ; jle ->overflow
            );
        }
    }

    #[inline]
    fn load_neg_offset(&mut self, offset: Count, proved: bool) {
        self.load_constant(offset);

        if self.checked && !proved {
            dynasm!(self.asm
                ; mov rcx, pointer
                ; sub rcx, mem_start
                ; cmp rcx, rax
                ; jl ->underflow
            );
        }
    }
}

impl JitCompilable for peephole::Program {
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R
    {
        k(self)
    }
}

impl<T: peephole::PeepholeCompilable + ?Sized> JitCompilable for T {
    fn with_peephole<F, R>(&self, k: F) -> R
        where F: FnOnce(&peephole::Program) -> R
    {
        k(&self.peephole_compile())
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

//! Rust helpers for Gecko's nsTArray.

use crate::gecko_bindings::bindings;
use crate::gecko_bindings::structs::{nsTArray, nsTArrayHeader, CopyableTArray};
use std::mem;
use std::ops::{Deref, DerefMut};
use std::slice;

impl<T> Deref for nsTArray<T> {
    type Target = [T];

    fn deref<'a>(&'a self) -> &'a [T] {
        unsafe { slice::from_raw_parts(self.slice_begin(), self.header().mLength as usize) }
    }
}

impl<T> DerefMut for nsTArray<T> {
    fn deref_mut<'a>(&'a mut self) -> &'a mut [T] {
        unsafe { slice::from_raw_parts_mut(self.slice_begin(), self.header().mLength as usize) }
    }
}

impl<T> nsTArray<T> {
    #[inline]
    fn header<'a>(&'a self) -> &'a nsTArrayHeader {
        debug_assert!(!self.mBuffer.is_null());
        unsafe { mem::transmute(self.mBuffer) }
    }

    // unsafe, since header may be in shared static or something
    unsafe fn header_mut<'a>(&'a mut self) -> &'a mut nsTArrayHeader {
        debug_assert!(!self.mBuffer.is_null());

        mem::transmute(self.mBuffer)
    }

    #[inline]
    unsafe fn slice_begin(&self) -> *mut T {
        debug_assert!(!self.mBuffer.is_null());
        (self.mBuffer as *const nsTArrayHeader).offset(1) as *mut _
    }

    /// Ensures the array has enough capacity at least to hold `cap` elements.
    ///
    /// NOTE: This doesn't call the constructor on the values!
    pub fn ensure_capacity(&mut self, cap: usize) {
        if cap >= self.len() {
            unsafe {
                bindings::Gecko_EnsureTArrayCapacity(
                    self as *mut nsTArray<T> as *mut _,
                    cap,
                    mem::size_of::<T>(),
                )
            }
        }
    }

    /// Clears the array storage without calling the destructor on the values.
    #[inline]
    pub unsafe fn clear(&mut self) {
        if self.len() != 0 {
            bindings::Gecko_ClearPODTArray(
                self as *mut nsTArray<T> as *mut _,
                mem::size_of::<T>(),
                mem::align_of::<T>(),
            );
        }
    }

    /// Clears a POD array. This is safe since copy types are memcopyable.
    #[inline]
    pub fn clear_pod(&mut self)
    where
        T: Copy,
    {
        unsafe { self.clear() }
    }

    /// Resize and set the length of the array to `len`.
    ///
    /// unsafe because this may leave the array with uninitialized elements.
    ///
    /// This will not call constructors.  If you need that, either manually add
    /// bindings or run the typed `EnsureCapacity` call on the gecko side.
    pub unsafe fn set_len(&mut self, len: u32) {
        // this can leak
        debug_assert!(len >= self.len() as u32);
        if self.len() == len as usize {
            return;
        }
        self.ensure_capacity(len as usize);
        self.header_mut().mLength = len;
    }

    /// Resizes an array containing only POD elements
    ///
    /// unsafe because this may leave the array with uninitialized elements.
    ///
    /// This will not leak since it only works on POD types (and thus doesn't assert)
    pub unsafe fn set_len_pod(&mut self, len: u32)
    where
        T: Copy,
    {
        if self.len() == len as usize {
            return;
        }
        self.ensure_capacity(len as usize);
        let header = self.header_mut();
        header.mLength = len;
    }

    /// Collects the given iterator into this array.
    ///
    /// Not unsafe because we won't leave uninitialized elements in the array.
    pub fn assign_from_iter_pod<I>(&mut self, iter: I)
    where
        T: Copy,
        I: ExactSizeIterator + Iterator<Item = T>,
    {
        debug_assert!(iter.len() <= 0xFFFFFFFF);
        unsafe {
            self.set_len_pod(iter.len() as u32);
        }
        self.iter_mut().zip(iter).for_each(|(r, v)| *r = v);
    }
}

impl<T> Deref for CopyableTArray<T> {
    type Target = nsTArray<T>;
    fn deref(&self) -> &Self::Target {
        &self._base
    }
}

impl<T> DerefMut for CopyableTArray<T> {
    fn deref_mut(&mut self) -> &mut nsTArray<T> {
        &mut self._base
    }
}
<commit_msg>style: Inline nsTArray::deref.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

//! Rust helpers for Gecko's nsTArray.

use crate::gecko_bindings::bindings;
use crate::gecko_bindings::structs::{nsTArray, nsTArrayHeader, CopyableTArray};
use std::mem;
use std::ops::{Deref, DerefMut};
use std::slice;

impl<T> Deref for nsTArray<T> {
    type Target = [T];

    #[inline]
    fn deref<'a>(&'a self) -> &'a [T] {
        unsafe { slice::from_raw_parts(self.slice_begin(), self.header().mLength as usize) }
    }
}

impl<T> DerefMut for nsTArray<T> {
    fn deref_mut<'a>(&'a mut self) -> &'a mut [T] {
        unsafe { slice::from_raw_parts_mut(self.slice_begin(), self.header().mLength as usize) }
    }
}

impl<T> nsTArray<T> {
    #[inline]
    fn header<'a>(&'a self) -> &'a nsTArrayHeader {
        debug_assert!(!self.mBuffer.is_null());
        unsafe { mem::transmute(self.mBuffer) }
    }

    // unsafe, since header may be in shared static or something
    unsafe fn header_mut<'a>(&'a mut self) -> &'a mut nsTArrayHeader {
        debug_assert!(!self.mBuffer.is_null());

        mem::transmute(self.mBuffer)
    }

    #[inline]
    unsafe fn slice_begin(&self) -> *mut T {
        debug_assert!(!self.mBuffer.is_null());
        (self.mBuffer as *const nsTArrayHeader).offset(1) as *mut _
    }

    /// Ensures the array has enough capacity at least to hold `cap` elements.
    ///
    /// NOTE: This doesn't call the constructor on the values!
    pub fn ensure_capacity(&mut self, cap: usize) {
        if cap >= self.len() {
            unsafe {
                bindings::Gecko_EnsureTArrayCapacity(
                    self as *mut nsTArray<T> as *mut _,
                    cap,
                    mem::size_of::<T>(),
                )
            }
        }
    }

    /// Clears the array storage without calling the destructor on the values.
    #[inline]
    pub unsafe fn clear(&mut self) {
        if self.len() != 0 {
            bindings::Gecko_ClearPODTArray(
                self as *mut nsTArray<T> as *mut _,
                mem::size_of::<T>(),
                mem::align_of::<T>(),
            );
        }
    }

    /// Clears a POD array. This is safe since copy types are memcopyable.
    #[inline]
    pub fn clear_pod(&mut self)
    where
        T: Copy,
    {
        unsafe { self.clear() }
    }

    /// Resize and set the length of the array to `len`.
    ///
    /// unsafe because this may leave the array with uninitialized elements.
    ///
    /// This will not call constructors.  If you need that, either manually add
    /// bindings or run the typed `EnsureCapacity` call on the gecko side.
    pub unsafe fn set_len(&mut self, len: u32) {
        // this can leak
        debug_assert!(len >= self.len() as u32);
        if self.len() == len as usize {
            return;
        }
        self.ensure_capacity(len as usize);
        self.header_mut().mLength = len;
    }

    /// Resizes an array containing only POD elements
    ///
    /// unsafe because this may leave the array with uninitialized elements.
    ///
    /// This will not leak since it only works on POD types (and thus doesn't assert)
    pub unsafe fn set_len_pod(&mut self, len: u32)
    where
        T: Copy,
    {
        if self.len() == len as usize {
            return;
        }
        self.ensure_capacity(len as usize);
        let header = self.header_mut();
        header.mLength = len;
    }

    /// Collects the given iterator into this array.
    ///
    /// Not unsafe because we won't leave uninitialized elements in the array.
    pub fn assign_from_iter_pod<I>(&mut self, iter: I)
    where
        T: Copy,
        I: ExactSizeIterator + Iterator<Item = T>,
    {
        debug_assert!(iter.len() <= 0xFFFFFFFF);
        unsafe {
            self.set_len_pod(iter.len() as u32);
        }
        self.iter_mut().zip(iter).for_each(|(r, v)| *r = v);
    }
}

impl<T> Deref for CopyableTArray<T> {
    type Target = nsTArray<T>;
    fn deref(&self) -> &Self::Target {
        &self._base
    }
}

impl<T> DerefMut for CopyableTArray<T> {
    fn deref_mut(&mut self) -> &mut nsTArray<T> {
        &mut self._base
    }
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::collections::{HashSet};
use std::iter::FromIterator;

extern crate xc2bit;
use self::xc2bit::*;

use *;
use objpool::*;

pub fn greedy_initial_placement(mcs: &[NetlistMacrocell]) -> Vec<[isize; MCS_PER_FB]> {
    let mut ret = Vec::new();

    // TODO: Number of FBs
    // FIXME: Hack for dedicated input
    for _ in 0..2 {
        ret.push([-1; MCS_PER_FB]);
    }
    if true {
        let x = ret.len();
        ret.push([-2; MCS_PER_FB]);
        ret[x][0] = -1;
    }

    // TODO: Handle LOCs
    let mut candidate_sites = Vec::new();
    if true {
        candidate_sites.push((2, 0));
    }
    for i in (0..2).rev() {
        for j in (0..MCS_PER_FB).rev() {
            candidate_sites.push((i, j));
        }
    }

    // Do the actual greedy assignment
    for i in 0..mcs.len() {
        if candidate_sites.len() == 0 {
            panic!("no more sites");
        }

        if let NetlistMacrocell::PinInputUnreg{..} = mcs[i] {} else {
            // Not an unregistered pin input. Special-case check for dedicated input
            if true && candidate_sites.len() == 1 {
                panic!("no more sites!");
            }
        }

        let (fb, mc) = candidate_sites.pop().unwrap();
        ret[fb][mc] = i as isize;
    }

    ret
}

fn compare_andterms(g: &NetlistGraph, a: ObjPoolIndex<NetlistGraphNode>, b: ObjPoolIndex<NetlistGraphNode>) -> bool {
    let a_ = g.nodes.get(a);
    let b_ = g.nodes.get(b);
    if let NetlistGraphNodeVariant::AndTerm{
        inputs_true: ref a_inp_true, inputs_comp: ref a_inp_comp, ..} = a_.variant {

        if let NetlistGraphNodeVariant::AndTerm{
            inputs_true: ref b_inp_true, inputs_comp: ref b_inp_comp, ..} = b_.variant {

            let mut a_inp_true_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut a_inp_comp_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut b_inp_true_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut b_inp_comp_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();

            for &x in a_inp_true {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    a_inp_true_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in a_inp_comp {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    a_inp_comp_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in b_inp_true {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    b_inp_true_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in b_inp_comp {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    b_inp_comp_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            a_inp_true_h == b_inp_true_h && a_inp_comp_h == b_inp_comp_h
        } else {
            panic!("not an and term");
        }
    } else {
        panic!("not an and term");
    }
}

pub enum AndTermAssignmentResult {
    Success([Option<ObjPoolIndex<NetlistGraphNode>>; ANDTERMS_PER_FB]),
    Failure(Vec<(usize, u32)>),
}

pub fn try_assign_andterms(g: &NetlistGraph, mcs: &[NetlistMacrocell], mc_assignment: &[[isize; MCS_PER_FB]])
    -> AndTermAssignmentResult {

    let mut ret = [None; ANDTERMS_PER_FB];

    AndTermAssignmentResult::Success(ret)
}
<commit_msg>xc2par: Fix sharing buried macrocells and inputs in initial assignment<commit_after>/*
Copyright (c) 2016-2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::collections::{HashSet};
use std::iter::FromIterator;

extern crate xc2bit;
use self::xc2bit::*;

use *;
use objpool::*;

// First element of tuple is anything, second element can only be pin input
pub fn greedy_initial_placement(mcs: &[NetlistMacrocell]) -> Vec<[(isize, isize); MCS_PER_FB]> {
    let mut ret = Vec::new();

    // TODO: Number of FBs
    // FIXME: Hack for dedicated input
    for _ in 0..2 {
        ret.push([(-1, -1); MCS_PER_FB]);
    }
    if true {
        let x = ret.len();
        ret.push([(-2, -2); MCS_PER_FB]);
        ret[x][0] = (-2, -1);
    }

    // TODO: Handle LOCs

    let mut candidate_sites_anything = Vec::new();
    let mut candidate_sites_pininput_conservative = Vec::new();
    let mut not_available_for_pininput_reg = HashSet::new();
    let mut not_available_for_pininput_unreg = HashSet::new();
    if true {
        candidate_sites_pininput_conservative.push((2, 0));
        not_available_for_pininput_reg.insert((2, 0));
    }
    for i in (0..2).rev() {
        for j in (0..MCS_PER_FB).rev() {
            candidate_sites_anything.push((i, j));
            candidate_sites_pininput_conservative.push((i, j));
        }
    }

    // Do the actual greedy assignment
    for i in 0..mcs.len() {
        match mcs[i] {
            NetlistMacrocell::PinOutput{..} => {
                if candidate_sites_anything.len() == 0 {
                    panic!("no more sites");
                }

                let (fb, mc) = candidate_sites_anything.pop().unwrap();
                // Remove from the pininput candidates as well. This requires that all PinOutput entries come first.
                candidate_sites_pininput_conservative.pop().unwrap();
                ret[fb][mc].0 = i as isize;
            },
            NetlistMacrocell::BuriedComb{..} => {
                if candidate_sites_anything.len() == 0 {
                    panic!("no more sites");
                }

                let (fb, mc) = candidate_sites_anything.pop().unwrap();
                // This is compatible with all input pin types
                ret[fb][mc].0 = i as isize;
            },
            NetlistMacrocell::BuriedReg{has_comb_fb, ..} => {
                if candidate_sites_anything.len() == 0 {
                    panic!("no more sites");
                }

                let (fb, mc) = candidate_sites_anything.pop().unwrap();
                // Cannot share with registered pin input
                not_available_for_pininput_reg.insert((fb, mc));
                if has_comb_fb {
                    // Not available to be shared at all
                    not_available_for_pininput_unreg.insert((fb, mc));
                }
                ret[fb][mc].0 = i as isize;
            },
            NetlistMacrocell::PinInputReg{..} => {
                // XXX This is not particularly efficient
                let mut idx = None;
                for j in (0..candidate_sites_pininput_conservative.len()).rev() {
                    if !not_available_for_pininput_reg.contains(&candidate_sites_pininput_conservative[j]) {
                        idx = Some(j);
                        break;
                    }
                }

                if idx.is_none() {
                    panic!("no more sites");
                }

                let (fb, mc) = candidate_sites_pininput_conservative.remove(idx.unwrap());
                ret[fb][mc].1 = i as isize;
            },
            NetlistMacrocell::PinInputUnreg{..} => {
                // XXX This is not particularly efficient
                let mut idx = None;
                for j in (0..candidate_sites_pininput_conservative.len()).rev() {
                    if !not_available_for_pininput_unreg.contains(&candidate_sites_pininput_conservative[j]) {
                        idx = Some(j);
                        break;
                    }
                }

                if idx.is_none() {
                    panic!("no more sites");
                }

                let (fb, mc) = candidate_sites_pininput_conservative.remove(idx.unwrap());
                ret[fb][mc].1 = i as isize;
            }
        }
    }

    ret
}

fn compare_andterms(g: &NetlistGraph, a: ObjPoolIndex<NetlistGraphNode>, b: ObjPoolIndex<NetlistGraphNode>) -> bool {
    let a_ = g.nodes.get(a);
    let b_ = g.nodes.get(b);
    if let NetlistGraphNodeVariant::AndTerm{
        inputs_true: ref a_inp_true, inputs_comp: ref a_inp_comp, ..} = a_.variant {

        if let NetlistGraphNodeVariant::AndTerm{
            inputs_true: ref b_inp_true, inputs_comp: ref b_inp_comp, ..} = b_.variant {

            let mut a_inp_true_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut a_inp_comp_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut b_inp_true_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();
            let mut b_inp_comp_h: HashSet<ObjPoolIndex<NetlistGraphNode>> = HashSet::new();

            for &x in a_inp_true {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    a_inp_true_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in a_inp_comp {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    a_inp_comp_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in b_inp_true {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    b_inp_true_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            for &x in b_inp_comp {
                let inp_net = g.nets.get(x);
                let src_node_idx = inp_net.source.unwrap().0;
                let src_node = g.nodes.get(src_node_idx);
                if let NetlistGraphNodeVariant::ZIADummyBuf{input, ..} = src_node.variant {
                    let zia_inp_net = g.nets.get(input);
                    let zia_src_node_idx = zia_inp_net.source.unwrap().0;
                    b_inp_comp_h.insert(zia_src_node_idx);
                } else {
                    panic!("mismatched node types");
                }
            }

            a_inp_true_h == b_inp_true_h && a_inp_comp_h == b_inp_comp_h
        } else {
            panic!("not an and term");
        }
    } else {
        panic!("not an and term");
    }
}

pub enum AndTermAssignmentResult {
    Success([Option<ObjPoolIndex<NetlistGraphNode>>; ANDTERMS_PER_FB]),
    Failure(Vec<(usize, u32)>),
}

pub fn try_assign_andterms(g: &NetlistGraph, mcs: &[NetlistMacrocell], mc_assignment: &[[(isize, isize); MCS_PER_FB]])
    -> AndTermAssignmentResult {

    let mut ret = [None; ANDTERMS_PER_FB];

    AndTermAssignmentResult::Success(ret)
}
<|endoftext|>


<commit_before>//! Global descriptor table

use core::convert::TryInto;
use core::mem;

use x86::segmentation::load_cs;
use x86::bits64::task::TaskStateSegment;
use x86::Ring;
use x86::dtables::{self, DescriptorTablePointer};
use x86::segmentation::{self, Descriptor as SegmentDescriptor, SegmentSelector};
use x86::task;

pub const GDT_NULL: usize = 0;
pub const GDT_KERNEL_CODE: usize = 1;
pub const GDT_KERNEL_DATA: usize = 2;
pub const GDT_KERNEL_KPCR: usize = 3;
pub const GDT_USER_CODE32_UNUSED: usize = 4;
pub const GDT_USER_DATA: usize = 5;
pub const GDT_USER_CODE: usize = 6;
pub const GDT_TSS: usize = 7;
pub const GDT_TSS_HIGH: usize = 8;
pub const GDT_CPU_ID_CONTAINER: usize = 9;

pub const GDT_A_PRESENT: u8 = 1 << 7;
pub const GDT_A_RING_0: u8 = 0 << 5;
pub const GDT_A_RING_1: u8 = 1 << 5;
pub const GDT_A_RING_2: u8 = 2 << 5;
pub const GDT_A_RING_3: u8 = 3 << 5;
pub const GDT_A_SYSTEM: u8 = 1 << 4;
pub const GDT_A_EXECUTABLE: u8 = 1 << 3;
pub const GDT_A_CONFORMING: u8 = 1 << 2;
pub const GDT_A_PRIVILEGE: u8 = 1 << 1;
pub const GDT_A_DIRTY: u8 = 1;

pub const GDT_A_TSS_AVAIL: u8 = 0x9;
pub const GDT_A_TSS_BUSY: u8 = 0xB;

pub const GDT_F_PAGE_SIZE: u8 = 1 << 7;
pub const GDT_F_PROTECTED_MODE: u8 = 1 << 6;
pub const GDT_F_LONG_MODE: u8 = 1 << 5;

static mut INIT_GDT: [GdtEntry; 4] = [
    // Null
    GdtEntry::new(0, 0, 0, 0),
    // Kernel code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel TLS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE)
];

#[thread_local]
pub static mut GDT: [GdtEntry; 10] = [
    // Null
    GdtEntry::new(0, 0, 0, 0),
    // Kernel code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel TLS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Dummy 32-bit user code - apparently necessary for SYSEXIT. We restrict it to ring 0 anyway.
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_PROTECTED_MODE),
    // User data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // User (64-bit) code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // TSS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_TSS_AVAIL, 0),
    // TSS must be 16 bytes long, twice the normal size
    GdtEntry::new(0, 0, 0, 0),
    // Unused entry which stores the CPU ID. This is necessary for paranoid interrupts as they have
    // no other way of determining it.
    GdtEntry::new(0, 0, 0, 0),
];

#[repr(C, align(16))]
pub struct ProcessorControlRegion {
    pub tcb_end: usize,
    pub user_rsp_tmp: usize,
    pub tss: TssWrapper,
}

// NOTE: Despite not using #[repr(packed)], we do know that while there may be some padding
// inserted before and after the TSS, the main TSS structure will remain intact.
#[repr(C, align(16))]
pub struct TssWrapper(pub TaskStateSegment);

#[thread_local]
pub static mut KPCR: ProcessorControlRegion = ProcessorControlRegion {
    tcb_end: 0,
    user_rsp_tmp: 0,
    tss: TssWrapper(TaskStateSegment {
        reserved: 0,
        rsp: [0; 3],
        reserved2: 0,
        ist: [0; 7],
        reserved3: 0,
        reserved4: 0,
        iomap_base: 0xFFFF
    }),
};

#[cfg(feature = "pti")]
pub unsafe fn set_tss_stack(stack: usize) {
    use super::pti::{PTI_CPU_STACK, PTI_CONTEXT_STACK};
    KPCR.tss.0.rsp[0] = (PTI_CPU_STACK.as_ptr() as usize + PTI_CPU_STACK.len()) as u64;
    PTI_CONTEXT_STACK = stack;
}

#[cfg(not(feature = "pti"))]
pub unsafe fn set_tss_stack(stack: usize) {
    KPCR.tss.0.rsp[0] = stack as u64;
}

// Initialize GDT
pub unsafe fn init() {
    {
        // Setup the initial GDT with TLS, so we can setup the TLS GDT (a little confusing)
        // This means that each CPU will have its own GDT, but we only need to define it once as a thread local

        let limit = (INIT_GDT.len() * mem::size_of::<GdtEntry>() - 1)
            .try_into()
            .expect("initial GDT way too large");
        let base = INIT_GDT.as_ptr() as *const SegmentDescriptor;

        let init_gdtr: DescriptorTablePointer<SegmentDescriptor> = DescriptorTablePointer {
            limit,
            base,
        };

        // Load the initial GDT, before we have access to thread locals
        dtables::lgdt(&init_gdtr);
    }

    // Load the segment descriptors
    load_cs(SegmentSelector::new(GDT_KERNEL_CODE as u16, Ring::Ring0));
    segmentation::load_ds(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_es(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_fs(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_gs(SegmentSelector::new(GDT_KERNEL_KPCR as u16, Ring::Ring0));
    segmentation::load_ss(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
}

/// Initialize GDT with TLS
pub unsafe fn init_paging(cpu_id: u32, tcb_offset: usize, stack_offset: usize) {
    // Set temporary TLS segment to the self-pointer of the Thread Control Block.
    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, tcb_offset as u64);

    // Now that we have access to thread locals, begin by getting a pointer to the Processor
    // Control Region.
    let kpcr = &mut KPCR;

    // Then, setup the AP's individual GDT
    let limit = (GDT.len() * mem::size_of::<GdtEntry>() - 1)
        .try_into()
        .expect("main GDT way too large");
    let base = GDT.as_ptr() as *const SegmentDescriptor;

    let gdtr: DescriptorTablePointer<SegmentDescriptor> = DescriptorTablePointer {
        limit,
        base,
    };

    // Once we have fetched the real KPCR address, set the TLS segment to the TCB pointer there.
    kpcr.tcb_end = (tcb_offset as *const usize).read();

    {
        // We can now access our TSS, via the KPCR, which is a thread local
        let tss = &kpcr.tss.0 as *const _ as usize as u64;
        let tss_lo = (tss & 0xFFFF_FFFF) as u32;
        let tss_hi = (tss >> 32) as u32;

        GDT[GDT_TSS].set_offset(tss_lo);
        GDT[GDT_TSS].set_limit(mem::size_of::<TaskStateSegment>() as u32);

        (&mut GDT[GDT_TSS_HIGH] as *mut GdtEntry).cast::<u32>().write(tss_hi);
    }

    // And finally, populate the last GDT entry with the current CPU ID, to allow paranoid
    // interrupt handlers to safely use TLS.
    (&mut GDT[GDT_CPU_ID_CONTAINER] as *mut GdtEntry).cast::<u32>().write(cpu_id);

    // Set the stack pointer to use when coming back from userspace.
    set_tss_stack(stack_offset);

    // Load the new GDT, which is correctly located in thread local storage.
    dtables::lgdt(&gdtr);

    // Ensure that GS always points to the KPCR in kernel space.
    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, kpcr as *mut _ as usize as u64);
    // Inside kernel space, GS should _always_ point to the TSS. When leaving userspace, `swapgs`
    // is called again, making the userspace GS always point to user data.
    x86::msr::wrmsr(x86::msr::IA32_KERNEL_GSBASE, 0);

    // Set the User TLS segment to zero, before we create any contexts and start scheduling.
    x86::msr::wrmsr(x86::msr::IA32_FS_BASE, 0);

    // Reload the segment descriptors
    load_cs(SegmentSelector::new(GDT_KERNEL_CODE as u16, Ring::Ring0));
    segmentation::load_ds(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_es(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_ss(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));

    // NOTE: FS has already been updated while calling set_tcb.
    // NOTE: We do not want to load GS again, since it has already been loaded into
    // GDT_KERNEL_KPCR. Instead, we use the base MSR to allow for a 64-bit offset.

    // Load the task register
    task::load_tr(SegmentSelector::new(GDT_TSS as u16, Ring::Ring0));

    let has_fsgsbase = raw_cpuid::CpuId::new()
        .get_extended_feature_info()
        .map_or(false, |extended_features| extended_features.has_fsgsbase());

    if cfg!(feature = "x86_fsgsbase") {
        assert!(has_fsgsbase, "running kernel with features not supported by the current CPU");
    }

    if has_fsgsbase {
        x86::controlregs::cr4_write(x86::controlregs::cr4() | x86::controlregs::Cr4::CR4_ENABLE_FSGSBASE);
    }
}

#[derive(Copy, Clone, Debug)]
#[repr(packed)]
pub struct GdtEntry {
    pub limitl: u16,
    pub offsetl: u16,
    pub offsetm: u8,
    pub access: u8,
    pub flags_limith: u8,
    pub offseth: u8
}

impl GdtEntry {
    pub const fn new(offset: u32, limit: u32, access: u8, flags: u8) -> Self {
        GdtEntry {
            limitl: limit as u16,
            offsetl: offset as u16,
            offsetm: (offset >> 16) as u8,
            access,
            flags_limith: flags & 0xF0 | ((limit >> 16) as u8) & 0x0F,
            offseth: (offset >> 24) as u8
        }
    }

    pub fn set_offset(&mut self, offset: u32) {
        self.offsetl = offset as u16;
        self.offsetm = (offset >> 16) as u8;
        self.offseth = (offset >> 24) as u8;
    }

    pub fn set_limit(&mut self, limit: u32) {
        self.limitl = limit as u16;
        self.flags_limith = self.flags_limith & 0xF0 | ((limit >> 16) as u8) & 0x0F;
    }
}
<commit_msg>Fix kernel TLS permissions in INIT_GDT<commit_after>//! Global descriptor table

use core::convert::TryInto;
use core::mem;

use x86::segmentation::load_cs;
use x86::bits64::task::TaskStateSegment;
use x86::Ring;
use x86::dtables::{self, DescriptorTablePointer};
use x86::segmentation::{self, Descriptor as SegmentDescriptor, SegmentSelector};
use x86::task;

pub const GDT_NULL: usize = 0;
pub const GDT_KERNEL_CODE: usize = 1;
pub const GDT_KERNEL_DATA: usize = 2;
pub const GDT_KERNEL_KPCR: usize = 3;
pub const GDT_USER_CODE32_UNUSED: usize = 4;
pub const GDT_USER_DATA: usize = 5;
pub const GDT_USER_CODE: usize = 6;
pub const GDT_TSS: usize = 7;
pub const GDT_TSS_HIGH: usize = 8;
pub const GDT_CPU_ID_CONTAINER: usize = 9;

pub const GDT_A_PRESENT: u8 = 1 << 7;
pub const GDT_A_RING_0: u8 = 0 << 5;
pub const GDT_A_RING_1: u8 = 1 << 5;
pub const GDT_A_RING_2: u8 = 2 << 5;
pub const GDT_A_RING_3: u8 = 3 << 5;
pub const GDT_A_SYSTEM: u8 = 1 << 4;
pub const GDT_A_EXECUTABLE: u8 = 1 << 3;
pub const GDT_A_CONFORMING: u8 = 1 << 2;
pub const GDT_A_PRIVILEGE: u8 = 1 << 1;
pub const GDT_A_DIRTY: u8 = 1;

pub const GDT_A_TSS_AVAIL: u8 = 0x9;
pub const GDT_A_TSS_BUSY: u8 = 0xB;

pub const GDT_F_PAGE_SIZE: u8 = 1 << 7;
pub const GDT_F_PROTECTED_MODE: u8 = 1 << 6;
pub const GDT_F_LONG_MODE: u8 = 1 << 5;

static mut INIT_GDT: [GdtEntry; 4] = [
    // Null
    GdtEntry::new(0, 0, 0, 0),
    // Kernel code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel TLS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE)
];

#[thread_local]
pub static mut GDT: [GdtEntry; 10] = [
    // Null
    GdtEntry::new(0, 0, 0, 0),
    // Kernel code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Kernel TLS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // Dummy 32-bit user code - apparently necessary for SYSEXIT. We restrict it to ring 0 anyway.
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_0 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_PROTECTED_MODE),
    // User data
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_SYSTEM | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // User (64-bit) code
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_SYSTEM | GDT_A_EXECUTABLE | GDT_A_PRIVILEGE, GDT_F_LONG_MODE),
    // TSS
    GdtEntry::new(0, 0, GDT_A_PRESENT | GDT_A_RING_3 | GDT_A_TSS_AVAIL, 0),
    // TSS must be 16 bytes long, twice the normal size
    GdtEntry::new(0, 0, 0, 0),
    // Unused entry which stores the CPU ID. This is necessary for paranoid interrupts as they have
    // no other way of determining it.
    GdtEntry::new(0, 0, 0, 0),
];

#[repr(C, align(16))]
pub struct ProcessorControlRegion {
    pub tcb_end: usize,
    pub user_rsp_tmp: usize,
    pub tss: TssWrapper,
}

// NOTE: Despite not using #[repr(packed)], we do know that while there may be some padding
// inserted before and after the TSS, the main TSS structure will remain intact.
#[repr(C, align(16))]
pub struct TssWrapper(pub TaskStateSegment);

#[thread_local]
pub static mut KPCR: ProcessorControlRegion = ProcessorControlRegion {
    tcb_end: 0,
    user_rsp_tmp: 0,
    tss: TssWrapper(TaskStateSegment {
        reserved: 0,
        rsp: [0; 3],
        reserved2: 0,
        ist: [0; 7],
        reserved3: 0,
        reserved4: 0,
        iomap_base: 0xFFFF
    }),
};

#[cfg(feature = "pti")]
pub unsafe fn set_tss_stack(stack: usize) {
    use super::pti::{PTI_CPU_STACK, PTI_CONTEXT_STACK};
    KPCR.tss.0.rsp[0] = (PTI_CPU_STACK.as_ptr() as usize + PTI_CPU_STACK.len()) as u64;
    PTI_CONTEXT_STACK = stack;
}

#[cfg(not(feature = "pti"))]
pub unsafe fn set_tss_stack(stack: usize) {
    KPCR.tss.0.rsp[0] = stack as u64;
}

// Initialize GDT
pub unsafe fn init() {
    {
        // Setup the initial GDT with TLS, so we can setup the TLS GDT (a little confusing)
        // This means that each CPU will have its own GDT, but we only need to define it once as a thread local

        let limit = (INIT_GDT.len() * mem::size_of::<GdtEntry>() - 1)
            .try_into()
            .expect("initial GDT way too large");
        let base = INIT_GDT.as_ptr() as *const SegmentDescriptor;

        let init_gdtr: DescriptorTablePointer<SegmentDescriptor> = DescriptorTablePointer {
            limit,
            base,
        };

        // Load the initial GDT, before we have access to thread locals
        dtables::lgdt(&init_gdtr);
    }

    // Load the segment descriptors
    load_cs(SegmentSelector::new(GDT_KERNEL_CODE as u16, Ring::Ring0));
    segmentation::load_ds(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_es(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_fs(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_gs(SegmentSelector::new(GDT_KERNEL_KPCR as u16, Ring::Ring0));
    segmentation::load_ss(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
}

/// Initialize GDT with TLS
pub unsafe fn init_paging(cpu_id: u32, tcb_offset: usize, stack_offset: usize) {
    // Set temporary TLS segment to the self-pointer of the Thread Control Block.
    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, tcb_offset as u64);

    // Now that we have access to thread locals, begin by getting a pointer to the Processor
    // Control Region.
    let kpcr = &mut KPCR;

    // Then, setup the AP's individual GDT
    let limit = (GDT.len() * mem::size_of::<GdtEntry>() - 1)
        .try_into()
        .expect("main GDT way too large");
    let base = GDT.as_ptr() as *const SegmentDescriptor;

    let gdtr: DescriptorTablePointer<SegmentDescriptor> = DescriptorTablePointer {
        limit,
        base,
    };

    // Once we have fetched the real KPCR address, set the TLS segment to the TCB pointer there.
    kpcr.tcb_end = (tcb_offset as *const usize).read();

    {
        // We can now access our TSS, via the KPCR, which is a thread local
        let tss = &kpcr.tss.0 as *const _ as usize as u64;
        let tss_lo = (tss & 0xFFFF_FFFF) as u32;
        let tss_hi = (tss >> 32) as u32;

        GDT[GDT_TSS].set_offset(tss_lo);
        GDT[GDT_TSS].set_limit(mem::size_of::<TaskStateSegment>() as u32);

        (&mut GDT[GDT_TSS_HIGH] as *mut GdtEntry).cast::<u32>().write(tss_hi);
    }

    // And finally, populate the last GDT entry with the current CPU ID, to allow paranoid
    // interrupt handlers to safely use TLS.
    (&mut GDT[GDT_CPU_ID_CONTAINER] as *mut GdtEntry).cast::<u32>().write(cpu_id);

    // Set the stack pointer to use when coming back from userspace.
    set_tss_stack(stack_offset);

    // Load the new GDT, which is correctly located in thread local storage.
    dtables::lgdt(&gdtr);

    // Ensure that GS always points to the KPCR in kernel space.
    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, kpcr as *mut _ as usize as u64);
    // Inside kernel space, GS should _always_ point to the TSS. When leaving userspace, `swapgs`
    // is called again, making the userspace GS always point to user data.
    x86::msr::wrmsr(x86::msr::IA32_KERNEL_GSBASE, 0);

    // Set the User TLS segment to zero, before we create any contexts and start scheduling.
    x86::msr::wrmsr(x86::msr::IA32_FS_BASE, 0);

    // Reload the segment descriptors
    load_cs(SegmentSelector::new(GDT_KERNEL_CODE as u16, Ring::Ring0));
    segmentation::load_ds(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_es(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));
    segmentation::load_ss(SegmentSelector::new(GDT_KERNEL_DATA as u16, Ring::Ring0));

    // NOTE: FS has already been updated while calling set_tcb.
    // NOTE: We do not want to load GS again, since it has already been loaded into
    // GDT_KERNEL_KPCR. Instead, we use the base MSR to allow for a 64-bit offset.

    // Load the task register
    task::load_tr(SegmentSelector::new(GDT_TSS as u16, Ring::Ring0));

    let has_fsgsbase = raw_cpuid::CpuId::new()
        .get_extended_feature_info()
        .map_or(false, |extended_features| extended_features.has_fsgsbase());

    if cfg!(feature = "x86_fsgsbase") {
        assert!(has_fsgsbase, "running kernel with features not supported by the current CPU");
    }

    if has_fsgsbase {
        x86::controlregs::cr4_write(x86::controlregs::cr4() | x86::controlregs::Cr4::CR4_ENABLE_FSGSBASE);
    }
}

#[derive(Copy, Clone, Debug)]
#[repr(packed)]
pub struct GdtEntry {
    pub limitl: u16,
    pub offsetl: u16,
    pub offsetm: u8,
    pub access: u8,
    pub flags_limith: u8,
    pub offseth: u8
}

impl GdtEntry {
    pub const fn new(offset: u32, limit: u32, access: u8, flags: u8) -> Self {
        GdtEntry {
            limitl: limit as u16,
            offsetl: offset as u16,
            offsetm: (offset >> 16) as u8,
            access,
            flags_limith: flags & 0xF0 | ((limit >> 16) as u8) & 0x0F,
            offseth: (offset >> 24) as u8
        }
    }

    pub fn set_offset(&mut self, offset: u32) {
        self.offsetl = offset as u16;
        self.offsetm = (offset >> 16) as u8;
        self.offseth = (offset >> 24) as u8;
    }

    pub fn set_limit(&mut self, limit: u32) {
        self.limitl = limit as u16;
        self.flags_limith = self.flags_limith & 0xF0 | ((limit >> 16) as u8) & 0x0F;
    }
}
<|endoftext|>


<commit_before>extern crate libc;
use std::ptr;
use std::mem;
use std::slice::from_raw_parts;
use std::string::String;
use std::collections::HashMap;
use std::boxed::Box;
use x11::xlib;
use x11::keysym;

use super::config::Config;
use super::workspace::{ Workspace, Workspaces };
use super::handler;

unsafe extern fn error_handler(display: *mut xlib::Display, event: *mut xlib::XErrorEvent) -> libc::c_int {
    // match event.error_code {
    //     xlib::BadAtom => {
    //         println!("bad atom");
    //     }
    //     _ => {}
    // }
    1
}

fn get_text_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) -> Option<String>{
    unsafe{
        let mut prop: xlib::XTextProperty = mem::zeroed();
        let r = xlib::XGetTextProperty(display, window, &mut prop, atom);
        if r == 0 {
            None
        }else{
            let s = String::from_raw_parts(prop.value, prop.nitems as usize, prop.nitems as usize + 1 ).clone();
            let text = Some(s);
            xlib::XFree(prop.value as *mut libc::c_void);
            text
        }
    }
}

fn get_window_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) {
    unsafe {
        let mut actual_type_return: u64 = 0;
        let mut actual_format_return: i32 = 0;
        let mut nitem_return: libc::c_ulong = 0;
        let mut bytes_after_return: libc::c_ulong = 0;
        let mut prop_return: *mut libc::c_uchar = mem::zeroed();

        let r = xlib::XGetWindowProperty(display, window, xlib::XA_WM_NAME,
                                         0, 0xFFFFFFFF, 0, 0,
                                         &mut actual_type_return,
                                         &mut actual_format_return,
                                         &mut nitem_return,
                                         &mut bytes_after_return,
                                         &mut prop_return);

        debug!("result get wp {}", r);
        if r == 0 {

        }else{
            debug!("actual format return {}", actual_format_return);
            if actual_format_return == 0 {

            }
            else {
                // let s = from_raw_parts(prop_return as *const libc::c_ulong, nitems_return as usize).iter()
                //     .map(|&c| c as u64)
                //     .collect();
                let s = String::from_raw_parts(prop_return, nitem_return as usize, nitem_return as usize + 1 );
                println!("{}", s);

            }
        }
    }
}

pub struct WindowManager {
    pub display: *mut xlib::Display,
    pub screen_num: libc::c_int,
    pub root:    xlib::Window,

    pub workspaces: Workspaces,
    config: Config
}

impl WindowManager {
    pub fn new() -> WindowManager {
	unsafe {
	    let display = xlib::XOpenDisplay(ptr::null());
            if display == ptr::null_mut() {
                panic!("can't open display")
            }
	    let screen_num  = xlib::XDefaultScreen(display);
	    let root    = xlib::XRootWindow(display, screen_num);

	    let mut wm = WindowManager {
	        display: display,
                screen_num: screen_num,
	        root:    root,

                config: Config::load(),
                workspaces: Workspaces::new()
            };
            wm.workspaces.create('1');
            wm
	}
    }

    pub fn run(&mut self) {
        let mask = xlib::SubstructureRedirectMask | xlib::SubstructureNotifyMask;
        let keymask = xlib::KeyPressMask | xlib::KeyReleaseMask;
        unsafe{
            xlib::XSetErrorHandler(Some(error_handler));
            xlib::XSelectInput(self.display, self.root,
                               mask | keymask);
            xlib::XSync(self.display, 0);
        }

        loop {
            //handle events here
            let mut e = xlib::XEvent{
                pad: [0; 24],
            };

            unsafe {
                xlib::XNextEvent(self.display, &mut e);
                let t = e.get_type();
                match t {
                    xlib::CreateNotify => {
                        let event: xlib::XCreateWindowEvent = From::from(e);

                        if event.override_redirect != 0 {
                            continue;
                        }

                        // change attributes before display
                        let mut attrs: xlib::XSetWindowAttributes = mem::zeroed();
                        attrs.event_mask = xlib::KeyReleaseMask;
                        let valuemask = xlib::CWEventMask;
                        xlib::XChangeWindowAttributes(self.display, event.window, valuemask, &mut attrs);

                        // add to workspace
                        // let mut w = self.workspaces.get_mut(&'1');
                        // match w {
                        //     Some(workspace) => {
                        //         workspace.add(event.window);
                        //         workspace.config(self);
                        //     }
                        //     None => {}
                        // }

                        // get window property
                        let n = get_text_property(self.display, event.window, xlib::XA_WM_NAME);
                        match n {
                            Some(s) => {debug!("create window {} for {}", event.window, s);}
                            None => {}
                        }
                    }
                    xlib::DestroyNotify => {
                        debug!("destroy notify");
                        let event: xlib::XDestroyWindowEvent = From::from(e);
                    }
                    xlib::MapNotify => {
                        debug!("map notify");
                        let event: xlib::XMapEvent = From::from(e);
                        println!("event {}", event.event);
                        println!("w {}", event.window);
                    }
                    xlib::UnmapNotify => {
                        debug!("unmap notify");
                        let event: xlib::XUnmapEvent = From::from(e);
                        let workspace = self.workspaces.current_workspace();
                        match workspace.contain(event.window) {
                            Some(index) => {
                                workspace.remove(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                            None => {}
                        }
                    }
                    xlib::MapRequest => {
                        debug!("map request");
                        let event: xlib::XMapRequestEvent = From::from(e);
                        println!("w {}", event.window);
                        xlib::XMapWindow(self.display, event.window);

                        // add app top-level window to workspace
                        let mut attrs: xlib::XWindowAttributes = mem::zeroed();
                        xlib::XGetWindowAttributes(self.display, event.window, &mut attrs);

                        let mut window_return: xlib::Window = 0;
                        let transientfor_hint = xlib::XGetTransientForHint(self.display, event.window, &mut window_return);
                        println!("trans hint {}", transientfor_hint);
                        if attrs.override_redirect == 0 && transientfor_hint == 0 {
                            debug!("top level window");
                            let mut workspace = self.workspaces.current();
                            if workspace.contain(event.window).is_none() {
                                workspace.add(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                        }

                    }
                    xlib::ClientMessage => {
                        debug!("client message");
                        // let event = cast_event::<xlib::XClientMessageEvent>(&e);
                        let event: xlib::XClientMessageEvent = From::from(e);
                    }
                    xlib::KeyRelease => {
                        debug!("key release");
                        let mut event: xlib::XKeyEvent = From::from(e);

                        if event.state > 0 {
                            let mut sym: xlib::KeySym = 0;
                            let status: *mut xlib::XComposeStatus = ptr::null_mut();
                            xlib::XLookupString(&mut event, ptr::null_mut(), 0, &mut sym, status);

                            let b = handler::KeyBind {
                                key: sym,
                                mask: event.state,
                            };
                            debug!("key {} {}", event.state, sym);
                            // let mut h: Box<handler::Handler>;
                            match self.config.bindsyms.get_mut(&b) {
                                Some(handler) => {
                                    handler.handle(&mut self.workspaces, self.display, self.screen_num);
                                }
                                None => {
                                    println!("no bind");
                                    continue;
                                }
                            }
                        }
                    }
                    xlib::ConfigureNotify => {
                        debug!("configure notify");
                    }
                    xlib::ConfigureRequest =>{
                        debug!("configure request");
                        // let mut change = xlib::XWindowChanges {
                        //     x: event.x,
                        //     y: event.y,
                        //     width: event.width,
                        //     height: event.height,
                        //     border_width: event.border_width,
                        //     sibling: event.above,
                        //     stack_mode: event.detail
                        // };
                        // debug!("config x: {}, y: {}, width: {}, height: {}",
                        //        change.x, change.y, change.width, change.height);
                        let mut event: xlib::XConfigureRequestEvent = From::from(e);
                        // xlib::XConfigureWindow(event.display, event.window, event.value_mask as u32, &mut change);

                    }
                    _ => {
                        debug!("unhandled event {}", t);
                    }
                }
            }
        }
    }
}
<commit_msg>client message<commit_after>extern crate libc;
use std::ptr;
use std::mem;
use std::ffi;
use std::str;
use std::slice::from_raw_parts;
use std::string::String;
use std::collections::HashMap;
use std::boxed::Box;
use x11::xlib;
use x11::keysym;

use super::config::Config;
use super::workspace::{ Workspace, Workspaces };
use super::handler;

unsafe extern fn error_handler(display: *mut xlib::Display, event: *mut xlib::XErrorEvent) -> libc::c_int {
    // match event.error_code {
    //     xlib::BadAtom => {
    //         println!("bad atom");
    //     }
    //     _ => {}
    // }
    1
}

fn get_atom_name(display: *mut xlib::Display, atom: xlib::Atom) -> Result<String, i32>{
    unsafe{
        let name = xlib::XGetAtomName(display, atom);
        if name == ptr::null_mut() {
            return Err(0);
        }
        let s = ffi::CStr::from_ptr(name);
        let s = s.to_bytes();
        match str::from_utf8(s) {
            Ok(v) => {
                let tmp = Ok(v.to_string());
                xlib::XFree(name as *mut libc::c_void);
                tmp
            }
            _ => Err(0)
        }
    }
}

fn get_text_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) -> Option<String>{
    unsafe{
        let mut prop: xlib::XTextProperty = mem::zeroed();
        let r = xlib::XGetTextProperty(display, window, &mut prop, atom);
        if r == 0 {
            None
        }else{
            let s = String::from_raw_parts(prop.value, prop.nitems as usize, prop.nitems as usize + 1 ).clone();
            let text = Some(s);
            xlib::XFree(prop.value as *mut libc::c_void);
            text
        }
    }
}

fn get_window_property(display: *mut xlib::Display, window: xlib::Window, atom: xlib::Atom) {
    unsafe {
        let mut actual_type_return: u64 = 0;
        let mut actual_format_return: i32 = 0;
        let mut nitem_return: libc::c_ulong = 0;
        let mut bytes_after_return: libc::c_ulong = 0;
        let mut prop_return: *mut libc::c_uchar = mem::zeroed();

        let r = xlib::XGetWindowProperty(display, window, xlib::XA_WM_NAME,
                                         0, 0xFFFFFFFF, 0, 0,
                                         &mut actual_type_return,
                                         &mut actual_format_return,
                                         &mut nitem_return,
                                         &mut bytes_after_return,
                                         &mut prop_return);

        debug!("result get wp {}", r);
        if r == 0 {

        }else{
            debug!("actual format return {}", actual_format_return);
            if actual_format_return == 0 {

            }
            else {
                // let s = from_raw_parts(prop_return as *const libc::c_ulong, nitems_return as usize).iter()
                //     .map(|&c| c as u64)
                //     .collect();
                let s = String::from_raw_parts(prop_return, nitem_return as usize, nitem_return as usize + 1 );
                println!("{}", s);

            }
        }
    }
}

pub struct WindowManager {
    pub display: *mut xlib::Display,
    pub screen_num: libc::c_int,
    pub root:    xlib::Window,

    pub workspaces: Workspaces,
    config: Config
}

impl WindowManager {
    pub fn new() -> WindowManager {
	unsafe {
	    let display = xlib::XOpenDisplay(ptr::null());
            if display == ptr::null_mut() {
                panic!("can't open display")
            }
	    let screen_num  = xlib::XDefaultScreen(display);
	    let root    = xlib::XRootWindow(display, screen_num);

	    let mut wm = WindowManager {
	        display: display,
                screen_num: screen_num,
	        root:    root,

                config: Config::load(),
                workspaces: Workspaces::new()
            };
            wm.workspaces.create('1');
            wm
	}
    }

    pub fn run(&mut self) {
        let mask = xlib::SubstructureRedirectMask | xlib::SubstructureNotifyMask;
        let keymask = xlib::KeyPressMask | xlib::KeyReleaseMask;
        unsafe{
            xlib::XSetErrorHandler(Some(error_handler));
            xlib::XSelectInput(self.display, self.root,
                               mask | keymask);
            xlib::XSync(self.display, 0);
        }

        loop {
            //handle events here
            let mut e = xlib::XEvent{
                pad: [0; 24],
            };

            unsafe {
                xlib::XNextEvent(self.display, &mut e);
                let t = e.get_type();
                match t {
                    xlib::CreateNotify => {
                        let event: xlib::XCreateWindowEvent = From::from(e);

                        if event.override_redirect != 0 {
                            continue;
                        }

                        // change attributes before display
                        let mut attrs: xlib::XSetWindowAttributes = mem::zeroed();
                        attrs.event_mask = xlib::KeyReleaseMask;
                        let valuemask = xlib::CWEventMask;
                        xlib::XChangeWindowAttributes(self.display, event.window, valuemask, &mut attrs);

                        // add to workspace
                        // let mut w = self.workspaces.get_mut(&'1');
                        // match w {
                        //     Some(workspace) => {
                        //         workspace.add(event.window);
                        //         workspace.config(self);
                        //     }
                        //     None => {}
                        // }

                        // get window property
                        let n = get_text_property(self.display, event.window, xlib::XA_WM_NAME);
                        match n {
                            Some(s) => {debug!("create window {} for {}", event.window, s);}
                            None => {}
                        }
                    }
                    xlib::DestroyNotify => {
                        debug!("destroy notify");
                        let event: xlib::XDestroyWindowEvent = From::from(e);
                    }
                    xlib::MapNotify => {
                        debug!("map notify");
                        let event: xlib::XMapEvent = From::from(e);
                        println!("event {}", event.event);
                        println!("w {}", event.window);
                    }
                    xlib::UnmapNotify => {
                        debug!("unmap notify");
                        let event: xlib::XUnmapEvent = From::from(e);
                        let workspace = self.workspaces.current_workspace();
                        match workspace.contain(event.window) {
                            Some(index) => {
                                workspace.remove(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                            None => {}
                        }
                    }
                    xlib::MapRequest => {
                        debug!("map request");
                        let event: xlib::XMapRequestEvent = From::from(e);
                        println!("w {}", event.window);
                        xlib::XMapWindow(self.display, event.window);

                        // add app top-level window to workspace
                        let mut attrs: xlib::XWindowAttributes = mem::zeroed();
                        xlib::XGetWindowAttributes(self.display, event.window, &mut attrs);

                        let mut window_return: xlib::Window = 0;
                        let transientfor_hint = xlib::XGetTransientForHint(self.display, event.window, &mut window_return);
                        println!("trans hint {}", transientfor_hint);
                        if attrs.override_redirect == 0 && transientfor_hint == 0 {
                            debug!("top level window");
                            let mut workspace = self.workspaces.current();
                            if workspace.contain(event.window).is_none() {
                                workspace.add(event.window);
                                workspace.config(self.display, self.screen_num);
                            }
                        }

                    }
                    xlib::ClientMessage => {
                        debug!("client message");
                        // let event = cast_event::<xlib::XClientMessageEvent>(&e);
                        let event: xlib::XClientMessageEvent = From::from(e);
                        let s = get_atom_name(self.display, event.message_type);
                        match s {
                            Ok(v) => println!("{}", v),
                            _ => {}
                        }

                        match event.format {
                            8 => {
                                for i in 0..20 {
                                    println!("{}", event.data.get_byte(i))
                                }
                            }
                            16 => {
                                for i in 0..10 {
                                    println!("{}", event.data.get_short(i))
                                }
                            }
                            32 => {
                                for i in 0..5 {
                                    let a = event.data.get_long(i);
                                    let s = get_atom_name(self.display, a as u64);
                                    match s {
                                        Ok(v) => println!("{}", v),
                                        _ => {}
                                    }
                                }
                            }
                            _ => {}
                        }
                    }
                    xlib::KeyRelease => {
                        debug!("key release");
                        let mut event: xlib::XKeyEvent = From::from(e);

                        if event.state > 0 {
                            let mut sym: xlib::KeySym = 0;
                            let status: *mut xlib::XComposeStatus = ptr::null_mut();
                            xlib::XLookupString(&mut event, ptr::null_mut(), 0, &mut sym, status);

                            let b = handler::KeyBind {
                                key: sym,
                                mask: event.state,
                            };
                            debug!("key {} {}", event.state, sym);
                            // let mut h: Box<handler::Handler>;
                            match self.config.bindsyms.get_mut(&b) {
                                Some(handler) => {
                                    handler.handle(&mut self.workspaces, self.display, self.screen_num);
                                }
                                None => {
                                    println!("no bind");
                                    continue;
                                }
                            }
                        }
                    }
                    xlib::ConfigureNotify => {
                        debug!("configure notify");
                    }
                    xlib::ConfigureRequest =>{
                        debug!("configure request");
                        // let mut change = xlib::XWindowChanges {
                        //     x: event.x,
                        //     y: event.y,
                        //     width: event.width,
                        //     height: event.height,
                        //     border_width: event.border_width,
                        //     sibling: event.above,
                        //     stack_mode: event.detail
                        // };
                        // debug!("config x: {}, y: {}, width: {}, height: {}",
                        //        change.x, change.y, change.width, change.height);
                        let mut event: xlib::XConfigureRequestEvent = From::from(e);
                        // xlib::XConfigureWindow(event.display, event.window, event.value_mask as u32, &mut change);

                    }
                    _ => {
                        debug!("unhandled event {}", t);
                    }
                }
            }
        }
    }
}
<|endoftext|>