<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use deriving::{path_local, path_std};
use deriving::generic::*;
use deriving::generic::ty::*;

use syntax::ast::{BinOpKind, Expr, MetaItem};
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use syntax_pos::Span;

pub fn expand_deriving_partial_eq(cx: &mut ExtCtxt,
                                  span: Span,
                                  mitem: &MetaItem,
                                  item: &Annotatable,
                                  push: &mut FnMut(Annotatable)) {
    // structures are equal if all fields are equal, and non equal, if
    // any fields are not equal or if the enum variants are different
    fn cs_eq(cx: &mut ExtCtxt, span: Span, substr: &Substructure) -> P<Expr> {
        cs_fold1(true, // use foldl
            |cx, span, subexpr, self_f, other_fs| {
                let other_f = match (other_fs.len(), other_fs.get(0)) {
                    (1, Some(o_f)) => o_f,
                    _ => cx.span_bug(span, "not exactly 2 arguments in `derive(PartialEq)`"),
                };

                let eq = cx.expr_binary(span, BinOpKind::Eq, self_f, other_f.clone());

                cx.expr_binary(span, BinOpKind::And, subexpr, eq)
            },
            |cx, args| {
                match args {
                    Some((span, self_f, other_fs)) => {
                        // Special-case the base case to generate cleaner code.
                        let other_f = match (other_fs.len(), other_fs.get(0)) {
                            (1, Some(o_f)) => o_f,
                            _ => {
                                cx.span_bug(span, "not exactly 2 arguments in `derive(PartialEq)`")
                            }
                        };

                        cx.expr_binary(span, BinOpKind::Eq, self_f, other_f.clone())
                    }
                    None => cx.expr_bool(span, true),
                }
            },
            Box::new(|cx, span, _, _| cx.expr_bool(span, false)),
            cx,
            span,
            substr)
    }
    fn cs_ne(cx: &mut ExtCtxt, span: Span, substr: &Substructure) -> P<Expr> {
        cs_fold1(true, // use foldl
            |cx, span, subexpr, self_f, other_fs| {
                let other_f = match (other_fs.len(), other_fs.get(0)) {
                    (1, Some(o_f)) => o_f,
                    _ => cx.span_bug(span, "not exactly 2 arguments in `derive(PartialEq)`"),
                };

                let eq = cx.expr_binary(span, BinOpKind::Ne, self_f, other_f.clone());

                cx.expr_binary(span, BinOpKind::Or, subexpr, eq)
            },
            |cx, args| {
                match args {
                    Some((span, self_f, other_fs)) => {
                        // Special-case the base case to generate cleaner code.
                        let other_f = match (other_fs.len(), other_fs.get(0)) {
                            (1, Some(o_f)) => o_f,
                            _ => {
                                cx.span_bug(span, "not exactly 2 arguments in `derive(PartialEq)`")
                            }
                        };

                        cx.expr_binary(span, BinOpKind::Ne, self_f, other_f.clone())
                    }
                    None => cx.expr_bool(span, false),
                }
            },
            Box::new(|cx, span, _, _| cx.expr_bool(span, true)),
            cx,
            span,
            substr)
    }

    macro_rules! md {
        ($name:expr, $f:ident) => { {
            let inline = cx.meta_word(span, Symbol::intern("inline"));
            let attrs = vec![cx.attribute(span, inline)];
            MethodDef {
                name: $name,
                generics: LifetimeBounds::empty(),
                explicit_self: borrowed_explicit_self(),
                args: vec![borrowed_self()],
                ret_ty: Literal(path_local!(bool)),
                attributes: attrs,
                is_unsafe: false,
                unify_fieldless_variants: true,
                combine_substructure: combine_substructure(Box::new(|a, b, c| {
                    $f(a, b, c)
                }))
            }
        } }
    }

    // avoid defining `ne` if we can
    // c-like enums, enums without any fields and structs without fields
    // can safely define only `eq`.
    let mut methods = vec![md!("eq", cs_eq)];
    if !is_type_without_fields(item) {
        methods.push(md!("ne", cs_ne));
    }

    let trait_def = TraitDef {
        span,
        attributes: Vec::new(),
        path: path_std!(cx, cmp::PartialEq),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        is_unsafe: false,
        supports_unions: false,
        methods,
        associated_types: Vec::new(),
    };
    trait_def.expand(cx, mitem, item, push)
}
<commit_msg>Abstract cs_eq for partial_eq<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use deriving::{path_local, path_std};
use deriving::generic::*;
use deriving::generic::ty::*;

use syntax::ast::{BinOpKind, Expr, MetaItem};
use syntax::ext::base::{Annotatable, ExtCtxt};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use syntax_pos::Span;

pub fn expand_deriving_partial_eq(cx: &mut ExtCtxt,
                                  span: Span,
                                  mitem: &MetaItem,
                                  item: &Annotatable,
                                  push: &mut FnMut(Annotatable)) {
    // structures are equal if all fields are equal, and non equal, if
    // any fields are not equal or if the enum variants are different
    fn cs_op(cx: &mut ExtCtxt,
             span: Span,
             substr: &Substructure,
             op: BinOpKind,
             combiner: BinOpKind,
             base: bool)
             -> P<Expr>
    {
        let op = |cx: &mut ExtCtxt, span: Span, self_f: P<Expr>, other_fs: &[P<Expr>]| {
            let other_f = match (other_fs.len(), other_fs.get(0)) {
                (1, Some(o_f)) => o_f,
                _ => cx.span_bug(span, "not exactly 2 arguments in `derive(PartialEq)`"),
            };

            cx.expr_binary(span, op, self_f, other_f.clone())
        };

        cs_fold1(true, // use foldl
            |cx, span, subexpr, self_f, other_fs| {
                let eq = op(cx, span, self_f, other_fs);
                cx.expr_binary(span, combiner, subexpr, eq)
            },
            |cx, args| {
                match args {
                    Some((span, self_f, other_fs)) => {
                        // Special-case the base case to generate cleaner code.
                        op(cx, span, self_f, other_fs)
                    }
                    None => cx.expr_bool(span, base),
                }
            },
            Box::new(|cx, span, _, _| cx.expr_bool(span, !base)),
            cx,
            span,
            substr)
    }

    fn cs_eq(cx: &mut ExtCtxt, span: Span, substr: &Substructure) -> P<Expr> {
        cs_op(cx, span, substr, BinOpKind::Eq, BinOpKind::And, true)
    }
    fn cs_ne(cx: &mut ExtCtxt, span: Span, substr: &Substructure) -> P<Expr> {
        cs_op(cx, span, substr, BinOpKind::Ne, BinOpKind::Or, false)
    }

    macro_rules! md {
        ($name:expr, $f:ident) => { {
            let inline = cx.meta_word(span, Symbol::intern("inline"));
            let attrs = vec![cx.attribute(span, inline)];
            MethodDef {
                name: $name,
                generics: LifetimeBounds::empty(),
                explicit_self: borrowed_explicit_self(),
                args: vec![borrowed_self()],
                ret_ty: Literal(path_local!(bool)),
                attributes: attrs,
                is_unsafe: false,
                unify_fieldless_variants: true,
                combine_substructure: combine_substructure(Box::new(|a, b, c| {
                    $f(a, b, c)
                }))
            }
        } }
    }

    // avoid defining `ne` if we can
    // c-like enums, enums without any fields and structs without fields
    // can safely define only `eq`.
    let mut methods = vec![md!("eq", cs_eq)];
    if !is_type_without_fields(item) {
        methods.push(md!("ne", cs_ne));
    }

    let trait_def = TraitDef {
        span,
        attributes: Vec::new(),
        path: path_std!(cx, cmp::PartialEq),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds::empty(),
        is_unsafe: false,
        supports_unions: false,
        methods,
        associated_types: Vec::new(),
    };
    trait_def.expand(cx, mitem, item, push)
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
use alloc::FrameAllocator;
use ::elf;
use memory::{Addr, PAGE_SIZE, PAddr, Page, PhysicalPage, VAddr, VirtualPage};

use core::marker::PhantomData;
use core::ops::{Index, IndexMut};
use core::{convert, fmt, intrinsics};

/// The number of entries in a page table.
pub const N_ENTRIES: usize = 512;
/// Size of a page table (in bytes)
pub const PAGE_TABLE_SIZE: usize = N_ENTRIES * PAGE_SIZE as usize;

/// Base virtual address of the PML4 table
pub const PML4_VADDR: u64 =  0xffffffff_fffff000;

/// A pointer to the PML4 table
pub const PML4_PTR: *mut Table<PML4Level> = PML4_VADDR as *mut _;

/// Mask to apply to a page table entry to isolate the flags
pub const ENTRY_FLAGS_MASK: u64 = (PAGE_SIZE as u64 - 1) as u64;

/// A page table
#[repr(C)]
pub struct Table<L>
where L: TableLevel { /// The entries in the page table.
                      entries: [Entry; N_ENTRIES]
                    , _level_marker: PhantomData<L>
                    }

impl<L:TableLevel> fmt::Debug for Table<L> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} at {:#p}"
              , unsafe { intrinsics::type_name::<L>() }
              , self)
    }
}

pub trait TableLevel {
    /// How much to shift an address by to find its index in this table.
    const SHIFT_AMOUNT: usize;
    const PAGE_SHIFT_AMOUNT: usize;

    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of(addr: VAddr) -> usize {
        (addr.as_usize() >> Self::SHIFT_AMOUNT) & 0b111111111
    }
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize;
}

pub enum PML4Level {}
pub enum PDPTLevel {}
pub enum PDLevel   {}
pub enum PTLevel   {}

impl TableLevel for PML4Level {
    const SHIFT_AMOUNT: usize = 39;
    const PAGE_SHIFT_AMOUNT: usize = 27;
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize {
        (page.number >> 27) & 0o777
    }
}
impl TableLevel for PDPTLevel {
    const SHIFT_AMOUNT: usize = 30;
    const PAGE_SHIFT_AMOUNT: usize = 18;
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize {
        (page.number >> 18) & 0o777
    }
}
impl TableLevel for PDLevel   {
    const SHIFT_AMOUNT: usize = 21;
    const PAGE_SHIFT_AMOUNT: usize = 9;
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize {
        (page.number >> 9) & 0o777
    }
}
impl TableLevel for PTLevel   {
    const SHIFT_AMOUNT: usize = 12;
    const PAGE_SHIFT_AMOUNT: usize = 0;
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize {
        (page.number >> 0) & 0o777
    }
}

pub trait Sublevel: TableLevel {
    type Next: TableLevel;
}
impl Sublevel for PML4Level {
    type Next = PDPTLevel;
}
impl Sublevel for PDPTLevel {
    type Next = PDLevel;
}
impl Sublevel for PDLevel {
    type Next = PTLevel;
}

impl Table<PML4Level> {
    #[inline]
    pub fn page_table_for(&self, page: VirtualPage) -> Option<&Table<PTLevel>> {
        self.next_table(page)
            .and_then(|pdpt| pdpt.next_table(page))
            .and_then(|pd| pd.next_table(page))
    }

    #[inline]
    pub fn page_table_mut_for(&mut self, page: VirtualPage)
                             -> Option<&mut Table<PTLevel>> {
        self.next_table_mut(page)
            .and_then(|pdpt| pdpt.next_table_mut(page))
            .and_then(|pd| pd.next_table_mut(page))
    }
}

impl<L: TableLevel> Index<usize> for Table<L> {
    type Output = Entry;

    fn index(&self, index: usize) -> &Entry {
        &self.entries[index]
    }
}

impl<L: TableLevel> IndexMut<usize> for Table<L> {
    fn index_mut(&mut self, index: usize) -> &mut Entry {
        &mut self.entries[index]
    }
}

impl<L: TableLevel> Index<VAddr> for Table<L> {
    type Output = Entry;

    #[inline] fn index(&self, addr: VAddr) -> &Entry {
        self.index(L::index_of(addr))
    }
}

impl<L: TableLevel> IndexMut<VAddr> for Table<L> {
    #[inline] fn index_mut(&mut self, addr: VAddr) -> &mut Entry {
        self.index_mut(L::index_of(addr))
    }
}
impl<L: TableLevel> Index<VirtualPage> for Table<L> {
    type Output = Entry;

    #[inline] fn index(&self, page: VirtualPage) -> &Entry {
        self.index(L::index_of_page(page))
    }
}

impl<L: TableLevel> IndexMut<VirtualPage> for Table<L> {
    #[inline] fn index_mut(&mut self, page: VirtualPage) -> &mut Entry {
        self.index_mut(L::index_of_page(page))
    }
}

impl<L: TableLevel> Table<L>  {

    /// Zeroes out the page table by setting all entries "unused"
    pub fn zero(&mut self) {
        trace!("zeroing {:?}", self);
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }

    /// Return the start physical address of this `Table`
    #[inline]
    pub fn start_paddr(&self) -> PAddr {
        PAddr::from(self as *const Self)
    }

    /// Return the `PhysicalPage` containing this table.
    #[inline]
    pub fn frame(&self) -> PhysicalPage {
        PhysicalPage::containing(self.start_paddr())
    }

}

impl<L: Sublevel> Table<L> {

    /// Returns the address of the next table, or None if none exists.
    #[inline]
    fn next_table_addr(&self, i: usize) -> Option<VAddr> {
        let flags = self[i].flags();
        if flags.contains(PRESENT) && !flags.contains(HUGE_PAGE) {
            let table_addr = self as *const _ as usize;
            Some(VAddr::from(table_addr << 9) | (i << 12))
        } else {
            None
        }
    }

    /// Returns the next table, or `None` if none exists
        #[inline]
    pub fn next_table(&self, i: VirtualPage) -> Option<&Table<L::Next>> {
        self.next_table_addr(L::index_of_page(i))
            .map(|table_addr| unsafe { &*(table_addr.as_ptr()) })
    }

    /// Mutably borrows the next table.
    #[inline]
    pub fn next_table_mut(&self, i: VirtualPage) -> Option<& mut Table<L::Next>> {
        trace!("{:?}, {:?}", self, i);
        self.next_table_addr(L::index_of_page(i))
            .map(|table_addr| unsafe { &mut *(table_addr.as_mut_ptr()) })
    }

    /// Returns the next table, creating it if it does not exist.
    pub fn create_next<A>(&mut self, i: VirtualPage, alloc: &mut A)
                         -> &mut Table<L::Next>
    where A: FrameAllocator {
        //println!("in create_next");
        if self.next_table(i).is_none() {
            assert!( !self[i].is_huge()
                   , "Couldn't create next table: huge pages not \
                      currently supported.");
            //print!("allocating...");
            let frame = unsafe {
                alloc.allocate()
                     // TODO: would we rather rewrite this to return
                     // a `Result`? I think so.
                     .expect("Couldn't map page, out of frames!")
            };
            //println!("done.");

            self[i].set(frame, PRESENT | WRITABLE);
            //println!("setted.");
            self.next_table_mut(i).unwrap().zero();
            trace!("zeroed");
        }
        self.next_table_mut(i).unwrap()
    }
}

bitflags! {
    pub flags EntryFlags: u64 {
        /// Present flag.
        /// Must be 1 to map a 2-MByte page or reference a page table.
        const PRESENT =         1 << 0,
        /// Writable flag.
        /// If 0, writes may not be allowed to the 2-MB region controlled
        /// by this entry
        const WRITABLE =        1 << 1
      , const USER_ACCESSIBLE = 1 << 2
      , const WRITE_THROUGH =   1 << 3
      , const NO_CACHE =        1 << 4
      , const ACCESSED =        1 << 5
      , const DIRTY =           1 << 6
      , const HUGE_PAGE =       1 << 7
      , const GLOBAL =          1 << 8
      , const NO_EXECUTE =      1 << 63
    }
}

impl EntryFlags {
    /// Returns true if this page is huge
    #[inline]
    pub fn is_huge(&self) -> bool {
        self.contains(HUGE_PAGE)
    }

    /// Returns true if this page is present
    #[inline]
    pub fn is_present(&self) -> bool {
        self.contains(PRESENT)
    }

    #[inline]
    pub fn set_present(&mut self, present: bool) -> &mut Self {
        if present { self.insert(PRESENT) }
        else { self.remove(PRESENT) }
        self
    }

    #[inline]
    pub fn set_writable(&mut self, writable: bool) -> &mut Self {
        if writable { self.insert(WRITABLE) }
        else { self.remove(WRITABLE) }
        self
    }

    #[inline]
    pub fn set_executable(&mut self, executable: bool) -> &mut Self {
        if executable { self.remove(NO_EXECUTE) }
        else { self.insert(NO_EXECUTE) }
        self
    }
}

#[derive(Debug)]
pub struct Entry(u64);

impl Entry {

    pub fn new(addr: PAddr) -> Self {
        assert!(addr.is_page_aligned());
        Entry(*addr)
    }

    // TODO: this is one of the worst names I have ever given a thing
    #[inline]
    pub fn do_huge(&self, offset: usize) -> Option<PhysicalPage> {
        if self.is_huge() {
            self.get_frame()
                .map(|start_frame| {
                    assert!( start_frame.number as usize % N_ENTRIES == 0
                           , "Start frame must be aligned on a 1GB boundary!");
                    start_frame + offset
                })
        } else {
            None
        }
    }

    /// Returns true if this is an unused entry
    #[inline]
    pub fn is_unused(&self) -> bool {
        self.0 == 0
    }

    /// Sets this entry to be unused
    #[inline(never)]
    pub fn set_unused(&mut self) {
        self.0 = 0;
    }

    /// Returns true if this page is huge
    #[inline]
    pub fn is_huge(&self) -> bool {
        self.flags().is_huge()
    }

    /// Access the entry's bitflags.
    #[inline]
    pub fn flags(&self) -> EntryFlags {
        EntryFlags::from_bits_truncate(self.0)
    }

    /// Returns the physical address pointed to by this page table entry
    #[inline]
    pub fn get_addr(&self) -> PAddr {
        PAddr::from(self.0 & PML4_VADDR)
    }

    /// Returns the frame in memory pointed to by this page table entry.
    pub fn get_frame(&self) -> Option<PhysicalPage> {
        if self.flags().is_present() {
            // If the entry is present, mask out bits 12-51 and
            Some(PhysicalPage::containing(self.get_addr()))
        } else {
            None
        }
    }

    pub fn set(&mut self, frame: PhysicalPage, flags: EntryFlags) {
        let addr: u64 = frame.base_addr().into();
        assert!(addr & !0x000fffff_fffff000 == 0);
        self.0 = addr | flags.bits();
    }

}

impl<'a> convert::From<&'a elf::Section<u64>> for EntryFlags {
    fn from(section: &'a elf::Section<u64>) -> Self {
        *EntryFlags::empty()
            .set_present(section.is_allocated())
            .set_writable(section.is_writable())
            .set_executable(section.is_executable())
    }
}
<commit_msg>refac(paging): start refactoring page table indexing<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
use alloc::FrameAllocator;
use ::elf;
use memory::{Addr, PAGE_SIZE, PAddr, Page, PhysicalPage, VAddr, VirtualPage};

use core::marker::PhantomData;
use core::ops::{Index, IndexMut};
use core::{convert, fmt, intrinsics};

/// The number of entries in a page table.
pub const N_ENTRIES: usize = 512;
/// Size of a page table (in bytes)
pub const PAGE_TABLE_SIZE: usize = N_ENTRIES * PAGE_SIZE as usize;

/// Base virtual address of the PML4 table
pub const PML4_VADDR: u64 =  0xffffffff_fffff000;

/// A pointer to the PML4 table
pub const PML4_PTR: *mut Table<PML4Level> = PML4_VADDR as *mut _;

/// Mask to apply to a page table entry to isolate the flags
pub const ENTRY_FLAGS_MASK: u64 = (PAGE_SIZE as u64 - 1) as u64;

/// A page table
#[repr(C)]
pub struct Table<L>
where L: TableLevel { /// The entries in the page table.
                      entries: [Entry; N_ENTRIES]
                    , _level_marker: PhantomData<L>
                    }

impl<L:TableLevel> fmt::Debug for Table<L> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} at {:#p}"
              , unsafe { intrinsics::type_name::<L>() }
              , self)
    }
}

pub trait TableLevel {
    /// How much to shift an address by to find its index in this table.
    const SHIFT_AMOUNT: usize;
    const PAGE_SHIFT_AMOUNT: usize;
    const PAGE_INDEX_MASK: usize = 0o777;

    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of(addr: VAddr) -> usize {
        (addr.as_usize() >> Self::SHIFT_AMOUNT) & 0b111111111
    }
    /// Returns the index in this table for the given virtual address
    #[inline]
    fn index_of_page(page: VirtualPage) -> usize {
        (page.number >> Self::PAGE_SHIFT_AMOUNT) & Self::PAGE_INDEX_MASK
    }
}

pub enum PML4Level {}
pub enum PDPTLevel {}
pub enum PDLevel   {}
pub enum PTLevel   {}

impl TableLevel for PML4Level {
    const SHIFT_AMOUNT: usize = 39;
    const PAGE_SHIFT_AMOUNT: usize = 27;
}
impl TableLevel for PDPTLevel {
    const SHIFT_AMOUNT: usize = 30;
    const PAGE_SHIFT_AMOUNT: usize = 18;
}
impl TableLevel for PDLevel   {
    const SHIFT_AMOUNT: usize = 21;
    const PAGE_SHIFT_AMOUNT: usize = 9;
}
impl TableLevel for PTLevel   {
    const SHIFT_AMOUNT: usize = 12;
    const PAGE_SHIFT_AMOUNT: usize = 0;
}

pub trait Sublevel: TableLevel {
    type Next: TableLevel;
}
impl Sublevel for PML4Level {
    type Next = PDPTLevel;
}
impl Sublevel for PDPTLevel {
    type Next = PDLevel;
}
impl Sublevel for PDLevel {
    type Next = PTLevel;
}

impl Table<PML4Level> {
    #[inline]
    pub fn page_table_for(&self, page: VirtualPage) -> Option<&Table<PTLevel>> {
        self.next_table(page)
            .and_then(|pdpt| pdpt.next_table(page))
            .and_then(|pd| pd.next_table(page))
    }

    #[inline]
    pub fn page_table_mut_for(&mut self, page: VirtualPage)
                             -> Option<&mut Table<PTLevel>> {
        self.next_table_mut(page)
            .and_then(|pdpt| pdpt.next_table_mut(page))
            .and_then(|pd| pd.next_table_mut(page))
    }
}

impl<L: TableLevel> Index<usize> for Table<L> {
    type Output = Entry;

    fn index(&self, index: usize) -> &Entry {
        &self.entries[index]
    }
}

impl<L: TableLevel> IndexMut<usize> for Table<L> {
    fn index_mut(&mut self, index: usize) -> &mut Entry {
        &mut self.entries[index]
    }
}

impl<L: TableLevel> Index<VAddr> for Table<L> {
    type Output = Entry;

    #[inline] fn index(&self, addr: VAddr) -> &Entry {
        self.index(L::index_of(addr))
    }
}

impl<L: TableLevel> IndexMut<VAddr> for Table<L> {
    #[inline] fn index_mut(&mut self, addr: VAddr) -> &mut Entry {
        self.index_mut(L::index_of(addr))
    }
}
impl<L: TableLevel> Index<VirtualPage> for Table<L> {
    type Output = Entry;

    #[inline] fn index(&self, page: VirtualPage) -> &Entry {
        self.index(L::index_of_page(page))
    }
}

impl<L: TableLevel> IndexMut<VirtualPage> for Table<L> {
    #[inline] fn index_mut(&mut self, page: VirtualPage) -> &mut Entry {
        self.index_mut(L::index_of_page(page))
    }
}

impl<L: TableLevel> Table<L>  {

    /// Zeroes out the page table by setting all entries "unused"
    pub fn zero(&mut self) {
        trace!("zeroing {:?}", self);
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }

    /// Return the start physical address of this `Table`
    #[inline]
    pub fn start_paddr(&self) -> PAddr {
        PAddr::from(self as *const Self)
    }

    /// Return the `PhysicalPage` containing this table.
    #[inline]
    pub fn frame(&self) -> PhysicalPage {
        PhysicalPage::containing(self.start_paddr())
    }

}

impl<L: Sublevel> Table<L> {

    /// Returns the address of the next table, or None if none exists.
    #[inline]
    fn next_table_addr(&self, i: usize) -> Option<VAddr> {
        let flags = self[i].flags();
        if flags.contains(PRESENT) && !flags.contains(HUGE_PAGE) {
            let table_addr = self as *const _ as usize;
            Some(VAddr::from(table_addr << 9) | (i << 12))
        } else {
            None
        }
    }

    /// Returns the next table, or `None` if none exists
        #[inline]
    pub fn next_table(&self, i: VirtualPage) -> Option<&Table<L::Next>> {
        self.next_table_addr(L::index_of_page(i))
            .map(|table_addr| unsafe { &*(table_addr.as_ptr()) })
    }

    /// Mutably borrows the next table.
    #[inline]
    pub fn next_table_mut(&self, i: VirtualPage) -> Option<& mut Table<L::Next>> {
        trace!("{:?}, {:?}", self, i);
        self.next_table_addr(L::index_of_page(i))
            .map(|table_addr| unsafe { &mut *(table_addr.as_mut_ptr()) })
    }

    /// Returns the next table, creating it if it does not exist.
    pub fn create_next<A>(&mut self, i: VirtualPage, alloc: &mut A)
                         -> &mut Table<L::Next>
    where A: FrameAllocator {
        //println!("in create_next");
        if self.next_table(i).is_none() {
            assert!( !self[i].is_huge()
                   , "Couldn't create next table: huge pages not \
                      currently supported.");
            //print!("allocating...");
            let frame = unsafe {
                alloc.allocate()
                     // TODO: would we rather rewrite this to return
                     // a `Result`? I think so.
                     .expect("Couldn't map page, out of frames!")
            };
            //println!("done.");

            self[i].set(frame, PRESENT | WRITABLE);
            //println!("setted.");
            self.next_table_mut(i).unwrap().zero();
            trace!("zeroed");
        }
        self.next_table_mut(i).unwrap()
    }
}

bitflags! {
    pub flags EntryFlags: u64 {
        /// Present flag.
        /// Must be 1 to map a 2-MByte page or reference a page table.
        const PRESENT =         1 << 0,
        /// Writable flag.
        /// If 0, writes may not be allowed to the 2-MB region controlled
        /// by this entry
        const WRITABLE =        1 << 1
      , const USER_ACCESSIBLE = 1 << 2
      , const WRITE_THROUGH =   1 << 3
      , const NO_CACHE =        1 << 4
      , const ACCESSED =        1 << 5
      , const DIRTY =           1 << 6
      , const HUGE_PAGE =       1 << 7
      , const GLOBAL =          1 << 8
      , const NO_EXECUTE =      1 << 63
    }
}

impl EntryFlags {
    /// Returns true if this page is huge
    #[inline]
    pub fn is_huge(&self) -> bool {
        self.contains(HUGE_PAGE)
    }

    /// Returns true if this page is present
    #[inline]
    pub fn is_present(&self) -> bool {
        self.contains(PRESENT)
    }

    #[inline]
    pub fn set_present(&mut self, present: bool) -> &mut Self {
        if present { self.insert(PRESENT) }
        else { self.remove(PRESENT) }
        self
    }

    #[inline]
    pub fn set_writable(&mut self, writable: bool) -> &mut Self {
        if writable { self.insert(WRITABLE) }
        else { self.remove(WRITABLE) }
        self
    }

    #[inline]
    pub fn set_executable(&mut self, executable: bool) -> &mut Self {
        if executable { self.remove(NO_EXECUTE) }
        else { self.insert(NO_EXECUTE) }
        self
    }
}

#[derive(Debug)]
pub struct Entry(u64);

impl Entry {

    pub fn new(addr: PAddr) -> Self {
        assert!(addr.is_page_aligned());
        Entry(*addr)
    }

    // TODO: this is one of the worst names I have ever given a thing
    #[inline]
    pub fn do_huge(&self, offset: usize) -> Option<PhysicalPage> {
        if self.is_huge() {
            self.get_frame()
                .map(|start_frame| {
                    assert!( start_frame.number as usize % N_ENTRIES == 0
                           , "Start frame must be aligned on a 1GB boundary!");
                    start_frame + offset
                })
        } else {
            None
        }
    }

    /// Returns true if this is an unused entry
    #[inline]
    pub fn is_unused(&self) -> bool {
        self.0 == 0
    }

    /// Sets this entry to be unused
    #[inline(never)]
    pub fn set_unused(&mut self) {
        self.0 = 0;
    }

    /// Returns true if this page is huge
    #[inline]
    pub fn is_huge(&self) -> bool {
        self.flags().is_huge()
    }

    /// Access the entry's bitflags.
    #[inline]
    pub fn flags(&self) -> EntryFlags {
        EntryFlags::from_bits_truncate(self.0)
    }

    /// Returns the physical address pointed to by this page table entry
    #[inline]
    pub fn get_addr(&self) -> PAddr {
        PAddr::from(self.0 & PML4_VADDR)
    }

    /// Returns the frame in memory pointed to by this page table entry.
    pub fn get_frame(&self) -> Option<PhysicalPage> {
        if self.flags().is_present() {
            // If the entry is present, mask out bits 12-51 and
            Some(PhysicalPage::containing(self.get_addr()))
        } else {
            None
        }
    }

    pub fn set(&mut self, frame: PhysicalPage, flags: EntryFlags) {
        let addr: u64 = frame.base_addr().into();
        assert!(addr & !0x000fffff_fffff000 == 0);
        self.0 = addr | flags.bits();
    }

}

impl<'a> convert::From<&'a elf::Section<u64>> for EntryFlags {
    fn from(section: &'a elf::Section<u64>) -> Self {
        *EntryFlags::empty()
            .set_present(section.is_allocated())
            .set_writable(section.is_writable())
            .set_executable(section.is_executable())
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#![recursion_limit = "512"]

extern crate proc_macro;
extern crate proc_macro2;
extern crate quote;
extern crate syn;

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
    braced, parenthesized,
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
    spanned::Spanned,
    token::Comma,
    ArgCaptured, Attribute, FnArg, Ident, Lit, LitBool, MetaNameValue, Pat, ReturnType, Token,
    Visibility,
};

struct Service {
    attrs: Vec<Attribute>,
    vis: Visibility,
    ident: Ident,
    rpcs: Vec<RpcMethod>,
}

struct RpcMethod {
    attrs: Vec<Attribute>,
    ident: Ident,
    args: Punctuated<ArgCaptured, Comma>,
    output: ReturnType,
}

impl Parse for Service {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis = input.parse()?;
        input.parse::<Token![trait]>()?;
        let ident = input.parse()?;
        let content;
        braced!(content in input);
        let mut rpcs = Vec::<RpcMethod>::new();
        while !content.is_empty() {
            rpcs.push(content.parse()?);
        }
        for rpc in &rpcs {
            if rpc.ident == "new" {
                return Err(syn::Error::new(
                    rpc.ident.span(),
                    format!(
                        "method name conflicts with generated fn `{}Client::new`",
                        ident
                    ),
                ));
            }
            if rpc.ident == "serve" {
                return Err(syn::Error::new(
                    rpc.ident.span(),
                    format!("method name conflicts with generated fn `{}::serve`", ident),
                ));
            }
        }
        Ok(Service {
            attrs,
            vis,
            ident,
            rpcs,
        })
    }
}

impl Parse for RpcMethod {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        input.parse::<Token![async]>()?;
        input.parse::<Token![fn]>()?;
        let ident = input.parse()?;
        let content;
        parenthesized!(content in input);
        let args: Punctuated<FnArg, Comma> = content.parse_terminated(FnArg::parse)?;
        let args = args
            .into_iter()
            .map(|arg| match arg {
                FnArg::Captured(captured) => match captured.pat {
                    Pat::Ident(_) => Ok(captured),
                    _ => {
                        return Err(syn::Error::new(
                            captured.pat.span(),
                            "patterns aren't allowed in RPC args",
                        ))
                    }
                },
                FnArg::SelfRef(self_ref) => {
                    return Err(syn::Error::new(
                        self_ref.span(),
                        "method args cannot start with self",
                    ))
                }
                FnArg::SelfValue(self_val) => {
                    return Err(syn::Error::new(
                        self_val.span(),
                        "method args cannot start with self",
                    ))
                }
                arg => {
                    return Err(syn::Error::new(
                        arg.span(),
                        "method args must be explicitly typed patterns",
                    ))
                }
            })
            .collect::<Result<_, _>>()?;
        let output = input.parse()?;
        input.parse::<Token![;]>()?;

        Ok(RpcMethod {
            attrs,
            ident,
            args,
            output,
        })
    }
}

// If `derive_serde` meta item is not present, defaults to cfg!(feature = "serde1").
// `derive_serde` can only be true when serde1 is enabled.
struct DeriveSerde(bool);

impl Parse for DeriveSerde {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        if input.is_empty() {
            return Ok(DeriveSerde(cfg!(feature = "serde1")));
        }
        match input.parse::<MetaNameValue>()? {
            MetaNameValue {
                ref ident, ref lit, ..
            } if ident == "derive_serde" => match lit {
                Lit::Bool(LitBool { value: true, .. }) if cfg!(feature = "serde1") => {
                    Ok(DeriveSerde(true))
                }
                Lit::Bool(LitBool { value: true, .. }) => Err(syn::Error::new(
                    lit.span(),
                    "To enable serde, first enable the `serde1` feature of tarpc",
                )),
                Lit::Bool(LitBool { value: false, .. }) => Ok(DeriveSerde(false)),
                lit => Err(syn::Error::new(
                    lit.span(),
                    "`derive_serde` expects a value of type `bool`",
                )),
            },
            MetaNameValue { ident, .. } => Err(syn::Error::new(
                ident.span(),
                "tarpc::service only supports one meta item, `derive_serde = {bool}`",
            )),
        }
    }
}

/// Generates:
/// - service trait
/// - serve fn
/// - client stub struct
/// - new_stub client factory fn
/// - Request and Response enums
/// - ResponseFut Future
#[proc_macro_attribute]
pub fn service(attr: TokenStream, input: TokenStream) -> TokenStream {
    let derive_serde = parse_macro_input!(attr as DeriveSerde);

    let Service {
        attrs,
        vis,
        ident,
        rpcs,
    } = parse_macro_input!(input as Service);
    let vis_repeated = std::iter::repeat(vis.clone());

    let camel_case_fn_names: Vec<String> = rpcs
        .iter()
        .map(|rpc| snake_to_camel(&rpc.ident.to_string()))
        .collect();
    let ref outputs: Vec<TokenStream2> = rpcs
        .iter()
        .map(|rpc| match rpc.output {
            ReturnType::Type(_, ref ty) => quote!(#ty),
            ReturnType::Default => quote!(()),
        })
        .collect();
    let future_types: Vec<Ident> = camel_case_fn_names
        .iter()
        .map(|name| Ident::new(&format!("{}Fut", name), ident.span()))
        .collect();
    let ref camel_case_idents: Vec<Ident> = rpcs
        .iter()
        .zip(camel_case_fn_names.iter())
        .map(|(rpc, name)| Ident::new(name, rpc.ident.span()))
        .collect();
    let camel_case_idents2 = camel_case_idents;

    let ref args: Vec<&Punctuated<ArgCaptured, Comma>> = rpcs.iter().map(|rpc| &rpc.args).collect();
    let ref arg_vars: Vec<Punctuated<&Pat, Comma>> = args
        .iter()
        .map(|args| args.iter().map(|arg| &arg.pat).collect())
        .collect();
    let arg_vars2 = arg_vars;
    let ref method_names: Vec<&Ident> = rpcs.iter().map(|rpc| &rpc.ident).collect();
    let method_attrs: Vec<_> = rpcs.iter().map(|rpc| &rpc.attrs).collect();

    let types_and_fns = rpcs
        .iter()
        .zip(future_types.iter())
        .zip(outputs.iter())
        .map(
            |(
                (
                    RpcMethod {
                        attrs, ident, args, ..
                    },
                    future_type,
                ),
                output,
            )| {
                let ty_doc = format!("The response future returned by {}.", ident);
                quote! {
                    #[doc = #ty_doc]
                    type #future_type: std::future::Future<Output = #output>;

                    #( #attrs )*
                    fn #ident(self, context: tarpc::context::Context, #args) -> Self::#future_type;
                }
            },
        );

    let service_name_repeated = std::iter::repeat(ident.clone());
    let service_name_repeated2 = service_name_repeated.clone();

    let client_ident = Ident::new(&format!("{}Client", ident), ident.span());
    let request_ident = Ident::new(&format!("{}Request", ident), ident.span());
    let request_ident_repeated = std::iter::repeat(request_ident.clone());
    let request_ident_repeated2 = request_ident_repeated.clone();
    let response_ident = Ident::new(&format!("{}Response", ident), ident.span());
    let response_ident_repeated = std::iter::repeat(response_ident.clone());
    let response_ident_repeated2 = response_ident_repeated.clone();
    let response_fut_name = format!("{}ResponseFut", ident);
    let response_fut_ident = Ident::new(&response_fut_name, ident.span());
    let response_fut_ident_repeated = std::iter::repeat(response_fut_ident.clone());
    let response_fut_ident_repeated2 = response_fut_ident_repeated.clone();
    let server_ident = Ident::new(&format!("Serve{}", ident), ident.span());

    let derive_serialize = if derive_serde.0 {
        quote!(#[derive(serde::Serialize, serde::Deserialize)])
    } else {
        quote!()
    };

    let tokens = quote! {
        #( #attrs )*
        #vis trait #ident: Clone {
            #( #types_and_fns )*

            /// Returns a serving function to use with tarpc::server::Server.
            fn serve(self) -> #server_ident<Self> {
                #server_ident { service: self }
            }
        }

        #[derive(Clone)]
        #vis struct #server_ident<S> {
            service: S,
        }

        impl<S> tarpc::server::Serve<#request_ident> for #server_ident<S>
            where S: #ident
        {
            type Resp = #response_ident;
            type Fut = #response_fut_ident<S>;

            fn serve(self, ctx: tarpc::context::Context, req: #request_ident) -> Self::Fut {
                match req {
                    #(
                        #request_ident_repeated::#camel_case_idents{ #arg_vars } => {
                            #response_fut_ident_repeated2::#camel_case_idents2(
                                #service_name_repeated2::#method_names(
                                    self.service, ctx, #arg_vars2))
                        }
                    )*
                }
            }
        }

        /// The request sent over the wire from the client to the server.
        #[derive(Debug)]
        #derive_serialize
        #vis enum #request_ident {
            #( #camel_case_idents{ #args } ),*
        }

        /// The response sent over the wire from the server to the client.
        #[derive(Debug)]
        #derive_serialize
        #vis enum #response_ident {
            #( #camel_case_idents(#outputs) ),*
        }

        /// A future resolving to a server response.
        #vis enum #response_fut_ident<S: #ident> {
            #( #camel_case_idents(<S as #service_name_repeated>::#future_types) ),*
        }

        impl<S: #ident> std::fmt::Debug for #response_fut_ident<S> {
            fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
                fmt.debug_struct(#response_fut_name).finish()
            }
        }

        impl<S: #ident> std::future::Future for #response_fut_ident<S> {
            type Output = #response_ident;

            fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>)
                -> std::task::Poll<#response_ident>
            {
                unsafe {
                    match std::pin::Pin::get_unchecked_mut(self) {
                        #(
                            #response_fut_ident_repeated::#camel_case_idents(resp) =>
                                std::pin::Pin::new_unchecked(resp)
                                    .poll(cx)
                                    .map(#response_ident_repeated::#camel_case_idents2),
                        )*
                    }
                }
            }
        }

        #[allow(unused)]
        #[derive(Clone, Debug)]
        /// The client stub that makes RPC calls to the server. Exposes a Future interface.
        #vis struct #client_ident<C = tarpc::client::Channel<#request_ident, #response_ident>>(C);

        impl<C> From<C> for #client_ident<C>
            where for <'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
        {
            fn from(client: C) -> Self {
                #client_ident(client)
            }
        }

        impl #client_ident {
            /// Returns a new client stub that sends requests over the given transport.
            #vis fn new<T>(config: tarpc::client::Config, transport: T)
                -> tarpc::client::NewClient<
                    Self,
                    tarpc::client::channel::RequestDispatch<#request_ident, #response_ident, T>>
            where
                T: tarpc::Transport<tarpc::ClientMessage<#request_ident>, tarpc::Response<#response_ident>>
            {
                let new_client = tarpc::client::new(config, transport);
                tarpc::client::NewClient {
                    client: #client_ident(new_client.client),
                    dispatch: new_client.dispatch,
                }
            }

        }

        impl<C> #client_ident<C>
            where for<'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
        {
            #(
                #[allow(unused)]
                #( #method_attrs )*
                #vis_repeated fn #method_names(&mut self, ctx: tarpc::context::Context, #args)
                    -> impl std::future::Future<Output = std::io::Result<#outputs>> + '_ {
                    let request = #request_ident_repeated2::#camel_case_idents { #arg_vars };
                    let resp = tarpc::Client::call(&mut self.0, ctx, request);
                    async move {
                        match resp.await? {
                            #response_ident_repeated2::#camel_case_idents2(msg) => std::result::Result::Ok(msg),
                            _ => unreachable!(),
                        }
                    }
                }
            )*
        }
    };

    tokens.into()
}

fn snake_to_camel(ident_str: &str) -> String {
    let mut camel_ty = String::new();
    let mut chars = ident_str.chars();

    let mut last_char_was_underscore = true;
    while let Some(c) = chars.next() {
        match c {
            '_' => last_char_was_underscore = true,
            c if last_char_was_underscore => {
                camel_ty.extend(c.to_uppercase());
                last_char_was_underscore = false;
            }
            c => camel_ty.extend(c.to_lowercase()),
        }
    }

    camel_ty
}

#[test]
fn snake_to_camel_basic() {
    assert_eq!(snake_to_camel("abc_def"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_suffix() {
    assert_eq!(snake_to_camel("abc_def_"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_prefix() {
    assert_eq!(snake_to_camel("_abc_def"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_consecutive() {
    assert_eq!(snake_to_camel("abc__def"), "AbcDef");
}

#[test]
fn snake_to_camel_capital_in_middle() {
    assert_eq!(snake_to_camel("aBc_dEf"), "AbcDef");
}
<commit_msg>Clippy fixes<commit_after>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#![recursion_limit = "512"]

extern crate proc_macro;
extern crate proc_macro2;
extern crate quote;
extern crate syn;

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
    braced, parenthesized,
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
    spanned::Spanned,
    token::Comma,
    ArgCaptured, Attribute, FnArg, Ident, Lit, LitBool, MetaNameValue, Pat, ReturnType, Token,
    Visibility,
};

struct Service {
    attrs: Vec<Attribute>,
    vis: Visibility,
    ident: Ident,
    rpcs: Vec<RpcMethod>,
}

struct RpcMethod {
    attrs: Vec<Attribute>,
    ident: Ident,
    args: Punctuated<ArgCaptured, Comma>,
    output: ReturnType,
}

impl Parse for Service {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis = input.parse()?;
        input.parse::<Token![trait]>()?;
        let ident = input.parse()?;
        let content;
        braced!(content in input);
        let mut rpcs = Vec::<RpcMethod>::new();
        while !content.is_empty() {
            rpcs.push(content.parse()?);
        }
        for rpc in &rpcs {
            if rpc.ident == "new" {
                return Err(syn::Error::new(
                    rpc.ident.span(),
                    format!(
                        "method name conflicts with generated fn `{}Client::new`",
                        ident
                    ),
                ));
            }
            if rpc.ident == "serve" {
                return Err(syn::Error::new(
                    rpc.ident.span(),
                    format!("method name conflicts with generated fn `{}::serve`", ident),
                ));
            }
        }
        Ok(Service {
            attrs,
            vis,
            ident,
            rpcs,
        })
    }
}

impl Parse for RpcMethod {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        input.parse::<Token![async]>()?;
        input.parse::<Token![fn]>()?;
        let ident = input.parse()?;
        let content;
        parenthesized!(content in input);
        let args: Punctuated<FnArg, Comma> = content.parse_terminated(FnArg::parse)?;
        let args = args
            .into_iter()
            .map(|arg| match arg {
                FnArg::Captured(captured) => match captured.pat {
                    Pat::Ident(_) => Ok(captured),
                    _ => Err(syn::Error::new(
                        captured.pat.span(),
                        "patterns aren't allowed in RPC args",
                    )),
                },
                FnArg::SelfRef(self_ref) => Err(syn::Error::new(
                    self_ref.span(),
                    "method args cannot start with self",
                )),
                FnArg::SelfValue(self_val) => Err(syn::Error::new(
                    self_val.span(),
                    "method args cannot start with self",
                )),
                arg => Err(syn::Error::new(
                    arg.span(),
                    "method args must be explicitly typed patterns",
                )),
            })
            .collect::<Result<_, _>>()?;
        let output = input.parse()?;
        input.parse::<Token![;]>()?;

        Ok(RpcMethod {
            attrs,
            ident,
            args,
            output,
        })
    }
}

// If `derive_serde` meta item is not present, defaults to cfg!(feature = "serde1").
// `derive_serde` can only be true when serde1 is enabled.
struct DeriveSerde(bool);

impl Parse for DeriveSerde {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        if input.is_empty() {
            return Ok(DeriveSerde(cfg!(feature = "serde1")));
        }
        match input.parse::<MetaNameValue>()? {
            MetaNameValue {
                ref ident, ref lit, ..
            } if ident == "derive_serde" => match lit {
                Lit::Bool(LitBool { value: true, .. }) if cfg!(feature = "serde1") => {
                    Ok(DeriveSerde(true))
                }
                Lit::Bool(LitBool { value: true, .. }) => Err(syn::Error::new(
                    lit.span(),
                    "To enable serde, first enable the `serde1` feature of tarpc",
                )),
                Lit::Bool(LitBool { value: false, .. }) => Ok(DeriveSerde(false)),
                lit => Err(syn::Error::new(
                    lit.span(),
                    "`derive_serde` expects a value of type `bool`",
                )),
            },
            MetaNameValue { ident, .. } => Err(syn::Error::new(
                ident.span(),
                "tarpc::service only supports one meta item, `derive_serde = {bool}`",
            )),
        }
    }
}

/// Generates:
/// - service trait
/// - serve fn
/// - client stub struct
/// - new_stub client factory fn
/// - Request and Response enums
/// - ResponseFut Future
#[proc_macro_attribute]
pub fn service(attr: TokenStream, input: TokenStream) -> TokenStream {
    let derive_serde = parse_macro_input!(attr as DeriveSerde);

    let Service {
        attrs,
        vis,
        ident,
        rpcs,
    } = parse_macro_input!(input as Service);
    let vis_repeated = std::iter::repeat(vis.clone());

    let camel_case_fn_names: Vec<String> = rpcs
        .iter()
        .map(|rpc| snake_to_camel(&rpc.ident.to_string()))
        .collect();
    let outputs: &Vec<TokenStream2> = &rpcs
        .iter()
        .map(|rpc| match rpc.output {
            ReturnType::Type(_, ref ty) => quote!(#ty),
            ReturnType::Default => quote!(()),
        })
        .collect();
    let future_types: Vec<Ident> = camel_case_fn_names
        .iter()
        .map(|name| Ident::new(&format!("{}Fut", name), ident.span()))
        .collect();
    let camel_case_idents: &Vec<Ident> = &rpcs
        .iter()
        .zip(camel_case_fn_names.iter())
        .map(|(rpc, name)| Ident::new(name, rpc.ident.span()))
        .collect();
    let camel_case_idents2 = camel_case_idents;

    let args: &Vec<&Punctuated<ArgCaptured, Comma>> = &rpcs.iter().map(|rpc| &rpc.args).collect();
    let arg_vars: &Vec<Punctuated<&Pat, Comma>> = &args
        .iter()
        .map(|args| args.iter().map(|arg| &arg.pat).collect())
        .collect();
    let arg_vars2 = arg_vars;
    let method_names: &Vec<&Ident> = &rpcs.iter().map(|rpc| &rpc.ident).collect();
    let method_attrs: Vec<_> = rpcs.iter().map(|rpc| &rpc.attrs).collect();

    let types_and_fns = rpcs
        .iter()
        .zip(future_types.iter())
        .zip(outputs.iter())
        .map(
            |(
                (
                    RpcMethod {
                        attrs, ident, args, ..
                    },
                    future_type,
                ),
                output,
            )| {
                let ty_doc = format!("The response future returned by {}.", ident);
                quote! {
                    #[doc = #ty_doc]
                    type #future_type: std::future::Future<Output = #output>;

                    #( #attrs )*
                    fn #ident(self, context: tarpc::context::Context, #args) -> Self::#future_type;
                }
            },
        );

    let service_name_repeated = std::iter::repeat(ident.clone());
    let service_name_repeated2 = service_name_repeated.clone();

    let client_ident = Ident::new(&format!("{}Client", ident), ident.span());
    let request_ident = Ident::new(&format!("{}Request", ident), ident.span());
    let request_ident_repeated = std::iter::repeat(request_ident.clone());
    let request_ident_repeated2 = request_ident_repeated.clone();
    let response_ident = Ident::new(&format!("{}Response", ident), ident.span());
    let response_ident_repeated = std::iter::repeat(response_ident.clone());
    let response_ident_repeated2 = response_ident_repeated.clone();
    let response_fut_name = format!("{}ResponseFut", ident);
    let response_fut_ident = Ident::new(&response_fut_name, ident.span());
    let response_fut_ident_repeated = std::iter::repeat(response_fut_ident.clone());
    let response_fut_ident_repeated2 = response_fut_ident_repeated.clone();
    let server_ident = Ident::new(&format!("Serve{}", ident), ident.span());

    let derive_serialize = if derive_serde.0 {
        quote!(#[derive(serde::Serialize, serde::Deserialize)])
    } else {
        quote!()
    };

    let tokens = quote! {
        #( #attrs )*
        #vis trait #ident: Clone {
            #( #types_and_fns )*

            /// Returns a serving function to use with tarpc::server::Server.
            fn serve(self) -> #server_ident<Self> {
                #server_ident { service: self }
            }
        }

        #[derive(Clone)]
        #vis struct #server_ident<S> {
            service: S,
        }

        impl<S> tarpc::server::Serve<#request_ident> for #server_ident<S>
            where S: #ident
        {
            type Resp = #response_ident;
            type Fut = #response_fut_ident<S>;

            fn serve(self, ctx: tarpc::context::Context, req: #request_ident) -> Self::Fut {
                match req {
                    #(
                        #request_ident_repeated::#camel_case_idents{ #arg_vars } => {
                            #response_fut_ident_repeated2::#camel_case_idents2(
                                #service_name_repeated2::#method_names(
                                    self.service, ctx, #arg_vars2))
                        }
                    )*
                }
            }
        }

        /// The request sent over the wire from the client to the server.
        #[derive(Debug)]
        #derive_serialize
        #vis enum #request_ident {
            #( #camel_case_idents{ #args } ),*
        }

        /// The response sent over the wire from the server to the client.
        #[derive(Debug)]
        #derive_serialize
        #vis enum #response_ident {
            #( #camel_case_idents(#outputs) ),*
        }

        /// A future resolving to a server response.
        #vis enum #response_fut_ident<S: #ident> {
            #( #camel_case_idents(<S as #service_name_repeated>::#future_types) ),*
        }

        impl<S: #ident> std::fmt::Debug for #response_fut_ident<S> {
            fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
                fmt.debug_struct(#response_fut_name).finish()
            }
        }

        impl<S: #ident> std::future::Future for #response_fut_ident<S> {
            type Output = #response_ident;

            fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>)
                -> std::task::Poll<#response_ident>
            {
                unsafe {
                    match std::pin::Pin::get_unchecked_mut(self) {
                        #(
                            #response_fut_ident_repeated::#camel_case_idents(resp) =>
                                std::pin::Pin::new_unchecked(resp)
                                    .poll(cx)
                                    .map(#response_ident_repeated::#camel_case_idents2),
                        )*
                    }
                }
            }
        }

        #[allow(unused)]
        #[derive(Clone, Debug)]
        /// The client stub that makes RPC calls to the server. Exposes a Future interface.
        #vis struct #client_ident<C = tarpc::client::Channel<#request_ident, #response_ident>>(C);

        impl<C> From<C> for #client_ident<C>
            where for <'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
        {
            fn from(client: C) -> Self {
                #client_ident(client)
            }
        }

        impl #client_ident {
            /// Returns a new client stub that sends requests over the given transport.
            #vis fn new<T>(config: tarpc::client::Config, transport: T)
                -> tarpc::client::NewClient<
                    Self,
                    tarpc::client::channel::RequestDispatch<#request_ident, #response_ident, T>>
            where
                T: tarpc::Transport<tarpc::ClientMessage<#request_ident>, tarpc::Response<#response_ident>>
            {
                let new_client = tarpc::client::new(config, transport);
                tarpc::client::NewClient {
                    client: #client_ident(new_client.client),
                    dispatch: new_client.dispatch,
                }
            }

        }

        impl<C> #client_ident<C>
            where for<'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
        {
            #(
                #[allow(unused)]
                #( #method_attrs )*
                #vis_repeated fn #method_names(&mut self, ctx: tarpc::context::Context, #args)
                    -> impl std::future::Future<Output = std::io::Result<#outputs>> + '_ {
                    let request = #request_ident_repeated2::#camel_case_idents { #arg_vars };
                    let resp = tarpc::Client::call(&mut self.0, ctx, request);
                    async move {
                        match resp.await? {
                            #response_ident_repeated2::#camel_case_idents2(msg) => std::result::Result::Ok(msg),
                            _ => unreachable!(),
                        }
                    }
                }
            )*
        }
    };

    tokens.into()
}

fn snake_to_camel(ident_str: &str) -> String {
    let mut camel_ty = String::new();
    let chars = ident_str.chars();

    let mut last_char_was_underscore = true;
    for c in chars {
        match c {
            '_' => last_char_was_underscore = true,
            c if last_char_was_underscore => {
                camel_ty.extend(c.to_uppercase());
                last_char_was_underscore = false;
            }
            c => camel_ty.extend(c.to_lowercase()),
        }
    }

    camel_ty
}

#[test]
fn snake_to_camel_basic() {
    assert_eq!(snake_to_camel("abc_def"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_suffix() {
    assert_eq!(snake_to_camel("abc_def_"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_prefix() {
    assert_eq!(snake_to_camel("_abc_def"), "AbcDef");
}

#[test]
fn snake_to_camel_underscore_consecutive() {
    assert_eq!(snake_to_camel("abc__def"), "AbcDef");
}

#[test]
fn snake_to_camel_capital_in_middle() {
    assert_eq!(snake_to_camel("aBc_dEf"), "AbcDef");
}
<|endoftext|>


<commit_before>use oppai_common::common::is_last_move_stupid;
use oppai_common::trajectory::{build_trajectories, build_trajectories_from, Trajectory};
use oppai_field::field::wave_diag;
use oppai_field::field::{Field, NonZeroPos, Pos};
use oppai_field::player::Player;
use std::iter;
use std::sync::atomic::{AtomicBool, Ordering};

fn mark_group(field: &Field, start_pos: Pos, player: Player, empty_board: &mut Vec<u32>) -> Vec<Pos> {
  let mut marks = Vec::new();
  wave_diag(field.width(), start_pos, |pos| {
    if empty_board[pos] != 0 {
      return false;
    }
    let cell = field.cell(pos);
    if cell.is_players_point(player) {
      empty_board[pos] = 1;
      marks.push(pos);
      true
    } else {
      false
    }
  });
  marks
}

fn collect_near_moves(field: &Field, player: Player, empty_board: &mut Vec<u32>) -> Vec<Pos> {
  let mut moves = Vec::new();
  for &pos in field
    .points_seq()
    .iter()
    .filter(|&&pos| field.cell(pos).is_players_point(player))
  {
    for &near_pos in field.directions(pos).iter() {
      if empty_board[near_pos] == 0 && field.is_putting_allowed(near_pos) {
        moves.push(near_pos);
        empty_board[near_pos] = 1;
      }
    }
  }
  for &pos in &moves {
    empty_board[pos] = 0;
  }
  moves
}

fn is_trajectoty_alive(field: &mut Field, trajectory: &Trajectory, player: Player, empty_board: &mut Vec<u32>) -> bool {
  if trajectory.points().iter().any(|&pos| !field.is_putting_allowed(pos)) {
    return false;
  }

  for &pos in trajectory.points() {
    field.put_point(pos, player);
  }

  let result = field.get_delta_score(player) > 0 || {
    let enemy = player.next();
    let moves = collect_near_moves(field, enemy, empty_board);
    moves.into_iter().any(|pos| {
      field.put_point(pos, player);
      let result = field.get_delta_score(player) > 0
        && trajectory
          .points()
          .iter()
          .all(|&trajectory_pos| field.cell(trajectory_pos).is_bound());
      field.undo();
      result && {
        let enemies_around = field
          .directions(pos)
          .iter()
          .filter(|&&pos| field.cell(pos).is_players_point(enemy))
          .count();
        enemies_around < 3
          || enemies_around == 3
            && field
              .directions(pos)
              .iter()
              .all(|&near_pos| !field.cell(near_pos).is_putting_allowed())
      }
    })
  };

  for _ in 0..trajectory.len() {
    field.undo();
  }

  result
}

fn is_trajectoty_viable(
  field: &mut Field,
  trajectory: &Trajectory,
  player: Player,
  empty_board: &mut Vec<u32>,
) -> bool {
  let moves = collect_near_moves(field, player, empty_board);
  let enemy = player.next();
  moves.into_iter().all(|enemy_pos| {
    field.put_point(enemy_pos, enemy);

    if field.get_delta_score(enemy) == 0 {
      field.undo();
      return true;
    }

    let result = is_trajectoty_alive(field, trajectory, player, empty_board);

    field.undo();

    result
  })
}

fn ladders_rec(
  field: &mut Field,
  player: Player,
  trajectory: &Trajectory,
  mut alpha: i32,
  beta: i32,
  empty_board: &mut Vec<u32>,
  should_stop: &AtomicBool,
  depth: u32,
) -> (Option<NonZeroPos>, i32, u32) {
  match *trajectory.points().as_slice() {
    [pos] => {
      field.put_point(pos, player);
      let cur_score = field.score(player);
      field.undo();
      (NonZeroPos::new(pos), cur_score, depth)
    }
    [pos1, pos2] => {
      let mut best_move = None;
      let mut capture_depth = 0;

      for &(our_pos, enemy_pos) in &[(pos1, pos2), (pos2, pos1)] {
        if trajectory.score() <= alpha || alpha >= beta {
          break;
        }

        if field.cell(our_pos).is_players_empty_base(player.next()) {
          continue;
        }

        if !field.put_point(our_pos, player) {
          panic!(
            "Failed to put a point to ({}, {}) on the field:\n{}",
            field.to_x(our_pos),
            field.to_y(our_pos),
            field,
          );
        }

        if is_last_move_stupid(field, our_pos, player) {
          field.undo();
          continue;
        }

        if field.get_delta_score(player) > 0 {
          let cur_score = field.score(player);
          field.undo();
          if cur_score > alpha {
            alpha = cur_score;
            best_move = NonZeroPos::new(our_pos);
            capture_depth = depth;
          }
          continue;
        }

        if !field.put_point(enemy_pos, player.next()) {
          panic!(
            "Failed to put a point to ({}, {}) on the field:\n{}",
            field.to_x(our_pos),
            field.to_y(our_pos),
            field,
          );
        }

        let trajectories = build_trajectories_from(field, our_pos, player, 2, empty_board, should_stop);

        if should_stop.load(Ordering::Relaxed) {
          break;
        }

        let marks = mark_group(field, our_pos, player, empty_board);

        for trajectory in trajectories {
          if alpha >= beta {
            break;
          }

          if trajectory.score() <= alpha {
            continue;
          }

          let (_, cur_score, cur_capture_depth) = ladders_rec(
            field,
            player,
            &trajectory,
            alpha,
            beta.min(trajectory.score()),
            empty_board,
            should_stop,
            depth + 1,
          );
          let cur_score = cur_score.min(trajectory.score());

          if cur_score > alpha && is_trajectoty_viable(field, &trajectory, player, empty_board) {
            alpha = cur_score;
            best_move = NonZeroPos::new(our_pos);
            capture_depth = cur_capture_depth;
          }
        }

        for pos in marks {
          empty_board[pos] = 0;
        }

        field.undo();
        field.undo();
      }

      (best_move, alpha, capture_depth)
    }
    _ => unreachable!("Trajectory with {} points", trajectory.len()),
  }
}

pub fn ladders(field: &mut Field, player: Player, should_stop: &AtomicBool) -> (Option<NonZeroPos>, i32, u32) {
  let mut empty_board = iter::repeat(0u32).take(field.length()).collect();

  let trajectories = build_trajectories(field, player, 2, &mut empty_board, &should_stop);

  info!("Solving ladders for {} trajectories.", trajectories.len());

  let mut alpha = field.score(player);
  let mut capture_depth = 0;
  let mut best_move = None;

  for trajectory in trajectories {
    if should_stop.load(Ordering::Relaxed) {
      break;
    }

    if trajectory.score() <= alpha {
      continue;
    }

    let marks = if let [pos1, _] = *trajectory.points().as_slice() {
      if let Some(&pos) = field
        .directions_diag(pos1)
        .iter()
        .find(|&&pos| field.cell(pos).is_players_point(player))
      {
        // mark one of near groups to not search trajectories from it
        mark_group(field, pos, player, &mut empty_board)
      } else {
        Vec::new()
      }
    } else {
      Vec::new()
    };

    let (cur_pos, cur_score, cur_capture_depth) = ladders_rec(
      field,
      player,
      &trajectory,
      alpha,
      trajectory.score(),
      &mut empty_board,
      should_stop,
      0,
    );
    let cur_score = cur_score.min(trajectory.score());
    if cur_score > alpha && is_trajectoty_viable(field, &trajectory, player, &mut empty_board) {
      alpha = cur_score;
      capture_depth = cur_capture_depth;
      best_move = cur_pos;
    }

    for pos in marks {
      empty_board[pos] = 0;
    }
  }

  (best_move, alpha, capture_depth)
}
<commit_msg>Fix preliminary termination for ladders.<commit_after>use oppai_common::common::is_last_move_stupid;
use oppai_common::trajectory::{build_trajectories, build_trajectories_from, Trajectory};
use oppai_field::field::wave_diag;
use oppai_field::field::{Field, NonZeroPos, Pos};
use oppai_field::player::Player;
use std::iter;
use std::sync::atomic::{AtomicBool, Ordering};

fn mark_group(field: &Field, start_pos: Pos, player: Player, empty_board: &mut Vec<u32>) -> Vec<Pos> {
  let mut marks = Vec::new();
  wave_diag(field.width(), start_pos, |pos| {
    if empty_board[pos] != 0 {
      return false;
    }
    let cell = field.cell(pos);
    if cell.is_players_point(player) {
      empty_board[pos] = 1;
      marks.push(pos);
      true
    } else {
      false
    }
  });
  marks
}

fn collect_near_moves(field: &Field, player: Player, empty_board: &mut Vec<u32>) -> Vec<Pos> {
  let mut moves = Vec::new();
  for &pos in field
    .points_seq()
    .iter()
    .filter(|&&pos| field.cell(pos).is_players_point(player))
  {
    for &near_pos in field.directions(pos).iter() {
      if empty_board[near_pos] == 0 && field.is_putting_allowed(near_pos) {
        moves.push(near_pos);
        empty_board[near_pos] = 1;
      }
    }
  }
  for &pos in &moves {
    empty_board[pos] = 0;
  }
  moves
}

fn is_trajectoty_alive(field: &mut Field, trajectory: &Trajectory, player: Player, empty_board: &mut Vec<u32>) -> bool {
  if trajectory.points().iter().any(|&pos| !field.is_putting_allowed(pos)) {
    return false;
  }

  for &pos in trajectory.points() {
    field.put_point(pos, player);
  }

  let result = field.get_delta_score(player) > 0 || {
    let enemy = player.next();
    let moves = collect_near_moves(field, enemy, empty_board);
    moves.into_iter().any(|pos| {
      field.put_point(pos, player);
      let result = field.get_delta_score(player) > 0
        && trajectory
          .points()
          .iter()
          .all(|&trajectory_pos| field.cell(trajectory_pos).is_bound());
      field.undo();
      result && {
        let enemies_around = field
          .directions(pos)
          .iter()
          .filter(|&&pos| field.cell(pos).is_players_point(enemy))
          .count();
        enemies_around < 3
          || enemies_around == 3
            && field
              .directions(pos)
              .iter()
              .all(|&near_pos| !field.cell(near_pos).is_putting_allowed())
      }
    })
  };

  for _ in 0..trajectory.len() {
    field.undo();
  }

  result
}

fn is_trajectoty_viable(
  field: &mut Field,
  trajectory: &Trajectory,
  player: Player,
  empty_board: &mut Vec<u32>,
) -> bool {
  let moves = collect_near_moves(field, player, empty_board);
  let enemy = player.next();
  moves.into_iter().all(|enemy_pos| {
    field.put_point(enemy_pos, enemy);

    if field.get_delta_score(enemy) == 0 {
      field.undo();
      return true;
    }

    let result = is_trajectoty_alive(field, trajectory, player, empty_board);

    field.undo();

    result
  })
}

fn ladders_rec(
  field: &mut Field,
  player: Player,
  trajectory: &Trajectory,
  mut alpha: i32,
  beta: i32,
  empty_board: &mut Vec<u32>,
  should_stop: &AtomicBool,
  depth: u32,
) -> (Option<NonZeroPos>, i32, u32) {
  match *trajectory.points().as_slice() {
    [pos] => {
      field.put_point(pos, player);
      let cur_score = field.score(player);
      field.undo();
      (NonZeroPos::new(pos), cur_score, depth)
    }
    [pos1, pos2] => {
      let mut best_move = None;
      let mut capture_depth = 0;

      for &(our_pos, enemy_pos) in &[(pos1, pos2), (pos2, pos1)] {
        if trajectory.score() <= alpha || alpha >= beta || should_stop.load(Ordering::Relaxed) {
          break;
        }

        if field.cell(our_pos).is_players_empty_base(player.next()) {
          continue;
        }

        if !field.put_point(our_pos, player) {
          panic!(
            "Failed to put a point to ({}, {}) on the field:\n{}",
            field.to_x(our_pos),
            field.to_y(our_pos),
            field,
          );
        }

        if is_last_move_stupid(field, our_pos, player) {
          field.undo();
          continue;
        }

        if field.get_delta_score(player) > 0 {
          let cur_score = field.score(player);
          field.undo();
          if cur_score > alpha {
            alpha = cur_score;
            best_move = NonZeroPos::new(our_pos);
            capture_depth = depth;
          }
          continue;
        }

        if !field.put_point(enemy_pos, player.next()) {
          panic!(
            "Failed to put a point to ({}, {}) on the field:\n{}",
            field.to_x(our_pos),
            field.to_y(our_pos),
            field,
          );
        }

        let trajectories = build_trajectories_from(field, our_pos, player, 2, empty_board, should_stop);

        if should_stop.load(Ordering::Relaxed) {
          field.undo();
          field.undo();
          break;
        }

        let marks = mark_group(field, our_pos, player, empty_board);

        for trajectory in trajectories {
          if alpha >= beta || should_stop.load(Ordering::Relaxed) {
            break;
          }

          if trajectory.score() <= alpha {
            continue;
          }

          let (_, cur_score, cur_capture_depth) = ladders_rec(
            field,
            player,
            &trajectory,
            alpha,
            beta.min(trajectory.score()),
            empty_board,
            should_stop,
            depth + 1,
          );
          let cur_score = cur_score.min(trajectory.score());

          if cur_score > alpha && is_trajectoty_viable(field, &trajectory, player, empty_board) {
            alpha = cur_score;
            best_move = NonZeroPos::new(our_pos);
            capture_depth = cur_capture_depth;
          }
        }

        for pos in marks {
          empty_board[pos] = 0;
        }

        field.undo();
        field.undo();
      }

      (best_move, alpha, capture_depth)
    }
    _ => unreachable!("Trajectory with {} points", trajectory.len()),
  }
}

pub fn ladders(field: &mut Field, player: Player, should_stop: &AtomicBool) -> (Option<NonZeroPos>, i32, u32) {
  let mut empty_board = iter::repeat(0u32).take(field.length()).collect();

  let trajectories = build_trajectories(field, player, 2, &mut empty_board, &should_stop);

  info!("Solving ladders for {} trajectories.", trajectories.len());

  let mut alpha = field.score(player);
  let mut capture_depth = 0;
  let mut best_move = None;

  for trajectory in trajectories {
    if should_stop.load(Ordering::Relaxed) {
      break;
    }

    if trajectory.score() <= alpha {
      continue;
    }

    let marks = if let [pos1, _] = *trajectory.points().as_slice() {
      if let Some(&pos) = field
        .directions_diag(pos1)
        .iter()
        .find(|&&pos| field.cell(pos).is_players_point(player))
      {
        // mark one of near groups to not search trajectories from it
        mark_group(field, pos, player, &mut empty_board)
      } else {
        Vec::new()
      }
    } else {
      Vec::new()
    };

    let (cur_pos, cur_score, cur_capture_depth) = ladders_rec(
      field,
      player,
      &trajectory,
      alpha,
      trajectory.score(),
      &mut empty_board,
      should_stop,
      0,
    );
    let cur_score = cur_score.min(trajectory.score());
    if cur_score > alpha && is_trajectoty_viable(field, &trajectory, player, &mut empty_board) {
      alpha = cur_score;
      capture_depth = cur_capture_depth;
      best_move = cur_pos;
    }

    for pos in marks {
      empty_board[pos] = 0;
    }
  }

  (best_move, alpha, capture_depth)
}
<|endoftext|>


<commit_before>use itertools::Itertools;
use pretty::{Arena, DocAllocator, DocBuilder};

use ast::{Expr, Literal, Pattern, SpannedExpr, SpannedPattern};
use types::{Type, pretty_print as pretty_type};

const INDENT: usize = 4;

fn forced_new_line<Id>(expr: &Expr<Id>) -> bool {
    match *expr {
        Expr::Block(..) |
        Expr::LetBindings(..) |
        Expr::Match(..) |
        Expr::TypeBindings(..) => true,
        _ => false,
    }
}

fn hang<'a, Id>(arena: &'a Arena<'a>, expr: &'a SpannedExpr<Id>) -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    let line = if forced_new_line(&expr.value) {
        arena.newline()
    } else {
        arena.space()
    };
    line.append(pretty_expr(arena, &expr.value))
}

pub fn pretty_pattern<'a, Id>(arena: &'a Arena<'a>,
                              pattern: &'a SpannedPattern<Id>)
                              -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    match pattern.value {
        Pattern::Constructor(ref ctor, ref args) => {
            chain![arena;
                ctor.as_ref(),
                arena.concat(args.iter().map(|arg| {
                    arena.text(" ").append(pretty_pattern(arena, arg))
                }))
            ]
        }
        Pattern::Ident(ref id) => arena.text(id.as_ref()),
        Pattern::Record { ref fields, ref types, .. } => {
            chain![arena;
                "{",
                arena.concat(types.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.0.as_ref(),
                        ","
                    ]
                }).chain(fields.iter().map(|field| {
                    match field.1 {
                        Some(ref new_name) => {
                            chain![arena;
                                arena.space(),
                                field.0.as_ref(),
                                " = ",
                                pretty_pattern(arena, new_name),
                                ","
                            ]
                        }
                        None => {
                            chain![arena;
                                arena.space(),
                                field.0.as_ref(),
                                ","
                            ]
                        }
                    }
                }))),
                if types.is_empty() && fields.is_empty() {
                    arena.nil()
                } else {
                    arena.space()
                },
                "}"
            ]
        }
        Pattern::Tuple { ref elems, .. } => {
            chain![arena;
                "(",
                arena.concat(elems
                    .iter()
                    .map(|elem| pretty_pattern(arena, elem))
                    .intersperse(arena.text(",").append(arena.space()))
                ),
                ")"
            ]
        }
    }
}

pub fn pretty_expr<'a, Id>(arena: &'a Arena<'a>, expr: &'a Expr<Id>) -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    let pretty = |expr: &'a SpannedExpr<_>| pretty_expr(arena, &expr.value);
    match *expr {
        Expr::App(ref func, ref args) => {
            pretty(func)
                .append(arena.concat(args.iter().map(|arg| arena.text(" ").append(pretty(arg)))))
        }
        Expr::Array(ref array) => {
            arena.text("[")
                .append(arena.concat(array.exprs.iter().map(|elem| pretty(elem).append(", "))))
                .append("]")
                .group()
        }
        Expr::Block(ref elems) => {
            arena.concat(elems.iter().map(|elem| pretty(elem).group().append(arena.newline())))
        }
        Expr::Ident(ref id) => arena.text(id.name.as_ref()),
        Expr::IfElse(ref body, ref if_true, ref if_false) => {
            chain![arena;
                "if",
                pretty(body),
                "then",
                pretty(if_true),
                "else",
                pretty(if_false)
            ]
        }
        Expr::Infix(ref l, ref op, ref r) => {
            chain![arena;
                pretty(l),
                " ",
                op.value.name.as_ref(),
                " ",
                pretty(r)
            ]
        }
        Expr::Lambda(ref lambda) => {
            chain![arena;
                "\\",
                arena.concat(lambda.args.iter().map(|arg| arena.text(arg.name.as_ref()).append(" "))),
                "->",
                hang(arena, &lambda.body)
            ]
        }
        Expr::LetBindings(ref binds, ref body) => {
            chain![arena;
                "let ",
                arena.concat(binds.iter().map(|bind| {
                    chain![arena;
                        chain![arena;
                            pretty_pattern(arena, &bind.name),
                            " ",
                            arena.concat(bind.args.iter().map(|arg| {
                                arena.text(arg.name.as_ref()).append(" ")
                            }))
                        ].group(),
                        match *bind.typ {
                            Type::Hole => arena.nil(),
                            ref typ => arena.text(": ").append(pretty_type(arena, typ)),
                        },
                        "=",
                        hang(arena, &bind.expr).nest(INDENT)
                    ].group().nest(INDENT)
                }).intersperse(arena.newline().append("and "))),
                arena.newline(),
                "in",
                arena.space(),
                pretty(body)
            ]
        }
        Expr::Literal(ref literal) => {
            match *literal {
                Literal::Byte(b) => arena.text(format!("b{}", b)),
                Literal::Char(c) => arena.text(format!("{:?}", c)),
                Literal::Float(f) => arena.text(format!("{}", f)),
                Literal::Int(i) => arena.text(format!("{}", i)),
                Literal::String(ref s) => arena.text(format!("{:?}", s)),
            }
        }
        Expr::Match(ref expr, ref alts) => {
            chain![arena;
                "match ",
                pretty(expr),
                " with",
                arena.newline(),
                arena.concat(alts.iter().map(|alt| {
                    chain![arena;
                        "| ",
                        pretty_pattern(arena, &alt.pattern),
                        " ->",
                        arena.space(),
                        pretty(&alt.expr).group().nest(INDENT)
                    ]
                }).intersperse(arena.newline()))
            ]
        }
        Expr::Projection(ref expr, ref field, _) => {
            chain![arena;
                pretty(expr),
                ".",
                field.as_ref()
            ]
        }
        Expr::Record { ref types, ref exprs, .. } => {
            chain![arena;
                "{",
                arena.concat(types.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.name.as_ref(),
                        ","
                    ]
                }).chain(exprs.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.name.as_ref(),
                        match field.value {
                            Some(ref expr) => {
                                chain![arena;
                                    " =",
                                    arena.space(),
                                    pretty(&expr)
                                ]
                            },
                            None => arena.nil(),
                        },
                        ","
                    ]
                })).map(|doc| doc.group().nest(INDENT))),
                if types.is_empty() && exprs.is_empty() {
                    arena.nil()
                } else {
                    arena.space()
                },
                "}"
            ]
        }
        Expr::Tuple{ ref elems, .. } => {
            arena.text("(")
                .append(arena.concat(elems.iter().map(|elem| pretty(elem).append(", "))))
                .append(")")
                .group()
        }
        Expr::TypeBindings(ref binds, ref body) => {
            chain![arena;
                "type ",
                arena.concat(binds.iter().map(|bind| {
                    chain![arena;
                        bind.name.as_ref(),
                        " ",
                        arena.concat(bind.alias.args.iter().map(|arg| {
                            arena.text(arg.id.as_ref()).append(" ")
                        })),
                        "=",
                        arena.space()
                            .append(pretty_type(arena, &bind.alias.unresolved_type()))
                            .group()
                    ].group().nest(INDENT)
                }).intersperse(arena.newline().append("and "))),
                arena.newline(),
                "in",
                arena.space(),
                pretty(body)
            ]
        }
        Expr::Error => arena.text("<error expr>"),
    }
}
<commit_msg>fix: Avoid writing `in`<commit_after>use itertools::Itertools;
use pretty::{Arena, DocAllocator, DocBuilder};

use ast::{Expr, Literal, Pattern, SpannedExpr, SpannedPattern};
use types::{Type, pretty_print as pretty_type};

const INDENT: usize = 4;

fn forced_new_line<Id>(expr: &Expr<Id>) -> bool {
    match *expr {
        Expr::Block(..) |
        Expr::LetBindings(..) |
        Expr::Match(..) |
        Expr::TypeBindings(..) => true,
        _ => false,
    }
}

fn hang<'a, Id>(arena: &'a Arena<'a>, expr: &'a SpannedExpr<Id>) -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    let line = if forced_new_line(&expr.value) {
        arena.newline()
    } else {
        arena.space()
    };
    line.append(pretty_expr(arena, &expr.value))
}

pub fn pretty_pattern<'a, Id>(arena: &'a Arena<'a>,
                              pattern: &'a SpannedPattern<Id>)
                              -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    match pattern.value {
        Pattern::Constructor(ref ctor, ref args) => {
            chain![arena;
                ctor.as_ref(),
                arena.concat(args.iter().map(|arg| {
                    arena.text(" ").append(pretty_pattern(arena, arg))
                }))
            ]
        }
        Pattern::Ident(ref id) => arena.text(id.as_ref()),
        Pattern::Record { ref fields, ref types, .. } => {
            chain![arena;
                "{",
                arena.concat(types.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.0.as_ref(),
                        ","
                    ]
                }).chain(fields.iter().map(|field| {
                    match field.1 {
                        Some(ref new_name) => {
                            chain![arena;
                                arena.space(),
                                field.0.as_ref(),
                                " = ",
                                pretty_pattern(arena, new_name),
                                ","
                            ]
                        }
                        None => {
                            chain![arena;
                                arena.space(),
                                field.0.as_ref(),
                                ","
                            ]
                        }
                    }
                }))),
                if types.is_empty() && fields.is_empty() {
                    arena.nil()
                } else {
                    arena.space()
                },
                "}"
            ]
        }
        Pattern::Tuple { ref elems, .. } => {
            chain![arena;
                "(",
                arena.concat(elems
                    .iter()
                    .map(|elem| pretty_pattern(arena, elem))
                    .intersperse(arena.text(",").append(arena.space()))
                ),
                ")"
            ]
        }
    }
}

pub fn pretty_expr<'a, Id>(arena: &'a Arena<'a>, expr: &'a Expr<Id>) -> DocBuilder<'a, Arena<'a>>
    where Id: AsRef<str>,
{
    let pretty = |expr: &'a SpannedExpr<_>| pretty_expr(arena, &expr.value);
    match *expr {
        Expr::App(ref func, ref args) => {
            pretty(func)
                .append(arena.concat(args.iter().map(|arg| arena.text(" ").append(pretty(arg)))))
        }
        Expr::Array(ref array) => {
            arena.text("[")
                .append(arena.concat(array.exprs.iter().map(|elem| pretty(elem).append(", "))))
                .append("]")
                .group()
        }
        Expr::Block(ref elems) => {
            arena.concat(elems.iter().map(|elem| pretty(elem).group().append(arena.newline())))
        }
        Expr::Ident(ref id) => arena.text(id.name.as_ref()),
        Expr::IfElse(ref body, ref if_true, ref if_false) => {
            chain![arena;
                "if",
                pretty(body),
                "then",
                pretty(if_true),
                "else",
                pretty(if_false)
            ]
        }
        Expr::Infix(ref l, ref op, ref r) => {
            chain![arena;
                pretty(l),
                " ",
                op.value.name.as_ref(),
                " ",
                pretty(r)
            ]
        }
        Expr::Lambda(ref lambda) => {
            chain![arena;
                "\\",
                arena.concat(lambda.args.iter().map(|arg| arena.text(arg.name.as_ref()).append(" "))),
                "->",
                hang(arena, &lambda.body)
            ]
        }
        Expr::LetBindings(ref binds, ref body) => {
            chain![arena;
                "let ",
                arena.concat(binds.iter().map(|bind| {
                    chain![arena;
                        chain![arena;
                            pretty_pattern(arena, &bind.name),
                            " ",
                            arena.concat(bind.args.iter().map(|arg| {
                                arena.text(arg.name.as_ref()).append(" ")
                            }))
                        ].group(),
                        match *bind.typ {
                            Type::Hole => arena.nil(),
                            ref typ => arena.text(": ").append(pretty_type(arena, typ)),
                        },
                        "=",
                        hang(arena, &bind.expr).nest(INDENT)
                    ].group().nest(INDENT)
                }).intersperse(arena.newline().append("and "))),
                arena.newline(),
                pretty(body)
            ]
        }
        Expr::Literal(ref literal) => {
            match *literal {
                Literal::Byte(b) => arena.text(format!("b{}", b)),
                Literal::Char(c) => arena.text(format!("{:?}", c)),
                Literal::Float(f) => arena.text(format!("{}", f)),
                Literal::Int(i) => arena.text(format!("{}", i)),
                Literal::String(ref s) => arena.text(format!("{:?}", s)),
            }
        }
        Expr::Match(ref expr, ref alts) => {
            chain![arena;
                "match ",
                pretty(expr),
                " with",
                arena.newline(),
                arena.concat(alts.iter().map(|alt| {
                    chain![arena;
                        "| ",
                        pretty_pattern(arena, &alt.pattern),
                        " ->",
                        arena.space(),
                        pretty(&alt.expr).group().nest(INDENT)
                    ]
                }).intersperse(arena.newline()))
            ]
        }
        Expr::Projection(ref expr, ref field, _) => {
            chain![arena;
                pretty(expr),
                ".",
                field.as_ref()
            ]
        }
        Expr::Record { ref types, ref exprs, .. } => {
            chain![arena;
                "{",
                arena.concat(types.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.name.as_ref(),
                        ","
                    ]
                }).chain(exprs.iter().map(|field| {
                    chain![arena;
                        arena.space(),
                        field.name.as_ref(),
                        match field.value {
                            Some(ref expr) => {
                                chain![arena;
                                    " =",
                                    arena.space(),
                                    pretty(&expr)
                                ]
                            },
                            None => arena.nil(),
                        },
                        ","
                    ]
                })).map(|doc| doc.group().nest(INDENT))),
                if types.is_empty() && exprs.is_empty() {
                    arena.nil()
                } else {
                    arena.space()
                },
                "}"
            ]
        }
        Expr::Tuple{ ref elems, .. } => {
            arena.text("(")
                .append(arena.concat(elems.iter().map(|elem| pretty(elem).append(", "))))
                .append(")")
                .group()
        }
        Expr::TypeBindings(ref binds, ref body) => {
            chain![arena;
                "type ",
                arena.concat(binds.iter().map(|bind| {
                    chain![arena;
                        bind.name.as_ref(),
                        " ",
                        arena.concat(bind.alias.args.iter().map(|arg| {
                            arena.text(arg.id.as_ref()).append(" ")
                        })),
                        "=",
                        arena.space()
                            .append(pretty_type(arena, &bind.alias.unresolved_type()))
                            .group()
                    ].group().nest(INDENT)
                }).intersperse(arena.newline().append("and "))),
                arena.newline(),
                pretty(body)
            ]
        }
        Expr::Error => arena.text("<error expr>"),
    }
}
<|endoftext|>