<commit_before>use super::{Key, Symbols, XkbFile, XkbKeySym};
use crate::{
    models::{DesktopModes, Layout},
    utils::UnwrapOrUnknownExt,
};
use snafu::{OptionExt, Snafu};

impl XkbFile {
    pub fn from_layout(name: &str, layout: Layout) -> Result<Self, Error> {
        let mut modes = vec![
            layout.modes.x11.clone().map(|x| ("x11", x)),
            layout.modes.win.clone().map(|x| ("win", x)),
            layout.modes.mac.clone().map(|x| ("mac", x)),
            layout.modes.chrome.clone().map(|x| ("chrome", x)),
        ]
        .into_iter()
        .flatten();

        let default = modes
            .next()
            .map(|(target, mode)| {
                Ok(Symbols {
                    id: "basic".to_string(),
                    name: format!("{} ({})", layout.name().unwrap_or_unknown(), target),
                    leading_includes: vec!["latin".to_string()],
                    keys: collect_keys(&mode, None)?,
                    trailing_includes: vec!["level3(ralt_switch)".to_string()],
                })
            })
            .context(NoXkbCompatibleModes {
                available_modes: layout.modes.available_modes(),
            })??;

        let others = modes
            .map(|(target, mode)| {
                Ok(Symbols {
                    id: target.to_string(),
                    name: format!("{} ({})", layout.name().unwrap_or_unknown(), target),
                    leading_includes: vec!["latin".to_string(), format!("{}(basic)", name)],
                    keys: collect_keys(&mode, Some(&default))?,
                    trailing_includes: vec!["level3(ralt_switch)".to_string()],
                })
            })
            .collect::<Result<Vec<Symbols>, Error>>()?;

        Ok(XkbFile { default, others })
    }
}

fn collect_keys(key_map: &DesktopModes, _default: Option<&Symbols>) -> Result<Vec<Key>, Error> {
    let default = key_map.get("default").cloned().context(NoDefaultKeyMap)?;
    let shift = key_map.get("shift").cloned().unwrap_or_default();
    let alt = key_map.get("alt").cloned().unwrap_or_default();
    let alt_shift = key_map.get("alt_shift").cloned().unwrap_or_default();

    let mut res = Vec::new();

    for (iso_code, default) in &*default {
        res.push(Key {
            iso_code: iso_code.to_string(),
            default: default.0.clone().map(|x| XkbKeySym(x)),
            shift: shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
            alt: alt.get_string(&iso_code).map(|x| XkbKeySym(x)),
            alt_shift: alt_shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
        });
    }

    Ok(res)
}

#[derive(Snafu, Debug)]
pub enum Error {
    #[snafu(display("No `default` keymap"))]
    NoDefaultKeyMap { backtrace: snafu::Backtrace },
    #[snafu(display("No XKB compatible modes, found: {}", available_modes.join(", ")))]
    NoXkbCompatibleModes {
        available_modes: Vec<String>,
        backtrace: snafu::Backtrace,
    },
}
<commit_msg>XKB: Add support for rendering dead key blocks<commit_after>use super::{Key, Symbols, XkbFile, XkbKeySym};
use crate::{
    models::{DesktopModes, Layout},
    utils::UnwrapOrUnknownExt,
};
use snafu::{OptionExt, Snafu};
use std::collections::BTreeMap;

impl XkbFile {
    pub fn from_layout(name: &str, layout: Layout) -> Result<Self, Error> {
        let mut modes = vec![
            layout.modes.x11.clone().map(|x| ("x11", x)),
            layout.modes.win.clone().map(|x| ("win", x)),
            layout.modes.mac.clone().map(|x| ("mac", x)),
            layout.modes.chrome.clone().map(|x| ("chrome", x)),
        ]
        .into_iter()
        .flatten();

        let default = modes
            .next()
            .map(|(target, mode)| {
                Ok(Symbols {
                    id: "basic".to_string(),
                    name: format!("{} ({})", layout.name().unwrap_or_unknown(), target),
                    leading_includes: vec!["latin".to_string()],
                    keys: collect_keys(&mode, None)?,
                    trailing_includes: vec!["level3(ralt_switch)".to_string()],
                })
            })
            .context(NoXkbCompatibleModes {
                available_modes: layout.modes.available_modes(),
            })??;

        let mut others = modes
            .map(|(target, mode)| {
                Ok(Symbols {
                    id: target.to_string(),
                    name: format!("{} ({})", layout.name().unwrap_or_unknown(), target),
                    leading_includes: vec![format!("{}(basic)", name)],
                    keys: collect_keys(&mode, Some(&default))?,
                    trailing_includes: vec!["level3(ralt_switch)".to_string()],
                })
            })
            .collect::<Result<Vec<Symbols>, Error>>()?;

        if let Some(dead_keys) = layout.dead_keys.as_ref() {
            for (target, mode_keys) in dead_keys {
                let parent = match target.as_str() {
                    "x11" => layout
                        .modes
                        .x11
                        .as_ref()
                        .context(DeadKeysForUnconfiguredTarget { target })?,
                    "win" => layout
                        .modes
                        .win
                        .as_ref()
                        .context(DeadKeysForUnconfiguredTarget { target })?,
                    "mac" => layout
                        .modes
                        .mac
                        .as_ref()
                        .context(DeadKeysForUnconfiguredTarget { target })?,
                    "chrome" => layout
                        .modes
                        .chrome
                        .as_ref()
                        .context(DeadKeysForUnconfiguredTarget { target })?,
                    _ => continue,
                };

                others.push(Symbols {
                    id: format!("{}_deadkeys", target.to_string()),
                    name: format!(
                        "{} ({}) (dead keys)",
                        layout.name().unwrap_or_unknown(),
                        target
                    ),
                    leading_includes: vec![format!("{}({})", name, target)],
                    keys: collect_dead_keys(&mode_keys, parent)?,
                    trailing_includes: vec![],
                });
            }
        }

        Ok(XkbFile { default, others })
    }
}

fn collect_keys(key_map: &DesktopModes, _default: Option<&Symbols>) -> Result<Vec<Key>, Error> {
    let default = key_map.get("default").cloned().context(NoDefaultKeyMap)?;
    let shift = key_map.get("shift").cloned().unwrap_or_default();
    let alt = key_map.get("alt").cloned().unwrap_or_default();
    let alt_shift = key_map.get("alt_shift").cloned().unwrap_or_default();

    let mut res = Vec::new();

    for (iso_code, default) in &*default {
        res.push(Key {
            iso_code: iso_code.to_string(),
            default: default.0.clone().map(|x| XkbKeySym(x)),
            shift: shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
            alt: alt.get_string(&iso_code).map(|x| XkbKeySym(x)),
            alt_shift: alt_shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
        });
    }

    Ok(res)
}

fn collect_dead_keys(
    key_map: &BTreeMap<String, Vec<String>>,
    parent: &DesktopModes,
) -> Result<Vec<Key>, Error> {
    fn char_to_dead(c: &str) -> Option<char> {
        let c = c.chars().next().expect("keysym can't be empty");
        let original = x11_keysymdef::lookup_by_codepoint(c)?;
        let name = original.names.get(0)?;
        let dead = format!("dead_{}", name);
        if let Some(record) = x11_keysymdef::lookup_by_name(&dead) {
            Some(record.unicode)
        } else {
            None
        }
    }

    let default = parent.get("default").cloned().context(NoDefaultKeyMap)?;
    let shift = parent.get("shift").cloned().unwrap_or_default();
    let alt = parent.get("alt").cloned().unwrap_or_default();
    let alt_shift = parent.get("alt_shift").cloned().unwrap_or_default();

    let mut overwritten_keys = BTreeMap::new();

    for (mode, map) in key_map {
        let parent_key = parent
            .get(mode)
            .expect("dead keys only defined for know modes");
        for key in map {
            let (iso_code, value) = parent_key
                .iter()
                .find(|(_iso, value)| value.0.as_ref() == Some(key))
                .expect("dead key not in parent mode");
            let value = value.0.as_ref().unwrap();

            let mut key = overwritten_keys.entry(iso_code).or_insert_with(|| Key {
                iso_code: iso_code.to_string(),
                default: default.get_string(&iso_code).map(|x| XkbKeySym(x)),
                shift: shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
                alt: alt.get_string(&iso_code).map(|x| XkbKeySym(x)),
                alt_shift: alt_shift.get_string(&iso_code).map(|x| XkbKeySym(x)),
            });
            let dead = char_to_dead(&value).map(|x| XkbKeySym(x.to_string()));
            let dead = if let Some(x) = dead {
                x
            } else {
                log::warn!(
                    "while generating XKB dead key block: `{}` cannot be turned into a dead key",
                    value
                );
                continue;
            };

            match mode.as_str() {
                "default" => key.default = Some(dead),
                "shift" => key.shift = Some(dead),
                "alt" => key.alt = Some(dead),
                "alt+shift" => key.alt_shift = Some(dead),
                _ => {}
            }
        }
    }

    Ok(overwritten_keys
        .into_iter()
        .map(|(_iso, val)| val)
        .collect())
}

#[derive(Snafu, Debug)]
pub enum Error {
    #[snafu(display("No `default` keymap"))]
    NoDefaultKeyMap { backtrace: snafu::Backtrace },
    #[snafu(display("No XKB compatible modes, found: {}", available_modes.join(", ")))]
    NoXkbCompatibleModes {
        available_modes: Vec<String>,
        backtrace: snafu::Backtrace,
    },
    #[snafu(display("Cannot set dead keys for unconfigured target `{}`", target))]
    DeadKeysForUnconfiguredTarget {
        target: String,
        backtrace: snafu::Backtrace,
    },
}
<|endoftext|>


<commit_before>use super::*;
use std::sync::atomic::fence;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Release;

pub enum InsertSearchResult {
    External,
    Internal(NodeCellRef),
    RightNode(NodeCellRef),
}

pub enum InsertToNodeResult {
    NoSplit,
    Split(NodeWriteGuard, NodeCellRef, Option<EntryKey>),
    SplitParentChanged,
}

pub enum RemoveSearchResult {
    External,
    RightNode(NodeCellRef),
    Internal(NodeCellRef),
}

pub enum SubNodeStatus {
    ExtNodeEmpty,
    InNodeEmpty,
    Relocate(usize, usize),
    Merge(usize, usize),
    Ok,
}

pub struct NodeSplit {
    pub new_right_node: NodeCellRef,
    pub left_node_latch: NodeWriteGuard,
    pub pivot: EntryKey,
    pub parent_latch: NodeWriteGuard,
}

pub struct RebalancingNodes {
    pub left_guard: NodeWriteGuard,
    pub left_ref: NodeCellRef,
    pub right_guard: NodeWriteGuard,
    pub right_right_guard: NodeWriteGuard, // for external pointer modification
    pub parent: NodeWriteGuard,
    pub parent_pos: usize,
}

pub struct RemoveResult {
    pub rebalancing: Option<RebalancingNodes>,
    pub empty: bool,
    pub removed: bool,
}

pub enum NodeData {
    External(Box<ExtNode>),
    Internal(Box<InNode>),
    Empty(Box<EmptyNode>),
    None,
}

pub struct EmptyNode {
    pub left: Option<NodeCellRef>,
    pub right: NodeCellRef,
}

impl NodeData {
    pub fn is_none(&self) -> bool {
        match self {
            &NodeData::None => true,
            _ => false,
        }
    }
    pub fn is_empty(&self) -> bool {
        match self {
            &NodeData::Empty(ref n) => true,
            &NodeData::External(ref n) => n.len == 0,
            &NodeData::Internal(ref n) => n.len == 0,
            &NodeData::None => unreachable!(),
        }
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        let len = self.len();
        if self.is_ext() {
            self.extnode().search(key)
        } else {
            self.innode().search(key)
        }
    }

    pub fn remove(&mut self, pos: usize) {
        match self {
            &mut NodeData::External(ref mut node) => node.remove_at(pos),
            &mut NodeData::Internal(ref mut node) => node.remove_at(pos),
            &mut NodeData::None | &mut NodeData::Empty(_) => unreachable!(),
        }
    }
    pub fn is_ext(&self) -> bool {
        match self {
            &NodeData::External(_) => true,
            &NodeData::Internal(_) => false,
            &NodeData::None | &NodeData::Empty(_) => panic!(),
        }
    }
    pub fn keys(&self) -> &[EntryKey] {
        if self.is_ext() {
            &self.extnode().keys
        } else {
            &self.innode().keys
        }
    }

    pub fn first_key(&self) -> &EntryKey {
        &self.keys()[0]
    }

    pub fn last_key(&self) -> &EntryKey {
        &self.keys()[self.len() - 1]
    }

    pub fn len(&self) -> usize {
        match self {
            &NodeData::Internal(ref node) => node.len,
            &NodeData::External(ref node) => node.len,
            &NodeData::Empty(_) => 0,
            &NodeData::None => unreachable!(),
        }
    }

    // check if the node will be half full after an item have been removed
    pub fn is_half_full(&self) -> bool {
        if self.is_none() {
            true
        } else {
            let len = self.len();
            if len == 0 {
                return false;
            }
            len >= NUM_KEYS / 2 && len > 1
        }
    }

    pub fn cannot_merge(&self) -> bool {
        self.len() > NUM_KEYS / 2
    }

    pub fn extnode_mut(&mut self) -> &mut ExtNode {
        match self {
            &mut NodeData::External(ref mut node) => node,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn innode_mut(&mut self) -> &mut InNode {
        match self {
            &mut NodeData::Internal(ref mut n) => n,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn extnode(&self) -> &ExtNode {
        match self {
            &NodeData::External(ref node) => node,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn ext_id(&self) -> Id {
        match self {
            &NodeData::External(ref node) => node.id,
            &NodeData::None => Id::unit_id(),
            &NodeData::Internal(_) | &NodeData::Empty(_) => unreachable!(self.type_name()),
        }
    }
    pub fn innode(&self) -> &InNode {
        match self {
            &NodeData::Internal(ref n) => n,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn type_name(&self) -> &'static str {
        match self {
            &NodeData::Internal(_) => "internal",
            &NodeData::External(_) => "external",
            &NodeData::None => "none",
            &NodeData::Empty(_) => "empty",
        }
    }
    pub fn key_at_right_node(&self, key: &EntryKey) -> Option<&NodeCellRef> {
        // return None;
        if self.len() > 0 {
            match self {
                &NodeData::Internal(ref n) => {
                    if &n.keys[n.len - 1] < key {
                        let right_node = n.right.read_unchecked();
                        if !right_node.is_none() {
                            let right_innode = right_node.innode();
                            if right_innode.keys.len() > 0 && &right_innode.keys[0] < key {
                                debug!(
                                    "found key to put to right internal page {:?}/{:?}",
                                    key, &right_innode.keys[0]
                                );
                                return Some(&n.right);
                            }
                        }
                    }
                }
                &NodeData::External(ref n) => {
                    if &n.keys[n.len - 1] < key {
                        let right_node = n.next.read_unchecked();
                        if !right_node.is_none() {
                            let right_extnode = right_node.extnode();
                            if right_extnode.keys.len() > 0 && &right_extnode.keys[0] < key {
                                debug!(
                                    "found key to put to right external page {:?}/{:?}",
                                    key, &right_extnode.keys[0]
                                );
                                return Some(&n.next);
                            }
                        }
                    }
                }
                &NodeData::Empty(ref n) => return Some(&n.right),
                _ => unreachable!(),
            }
        }
        return None;
    }

    pub fn left_ref_mut(&mut self) -> Option<&mut NodeCellRef> {
        match self {
            &mut NodeData::External(ref mut n) => Some(&mut n.prev),
            _ => None,
        }
    }

    pub fn left_ref(&self) -> Option<&NodeCellRef> {
        match self {
            &NodeData::External(ref n) => Some(&n.prev),
            _ => None,
        }
    }

    pub fn right_ref_mut(&mut self) -> Option<&mut NodeCellRef> {
        match self {
            &mut NodeData::External(ref mut n) => Some(&mut n.next),
            &mut NodeData::Internal(ref mut n) => Some(&mut n.right),
            &mut NodeData::Empty(ref mut n) => Some(&mut n.right),
            &mut NodeData::None => None,
        }
    }

    pub fn right_ref(&self) -> Option<&NodeCellRef> {
        match self {
            &NodeData::External(ref n) => Some(&n.next),
            &NodeData::Internal(ref n) => Some(&n.right),
            &NodeData::Empty(ref n) => Some(&n.right),
            &NodeData::None => None,
        }
    }

    pub fn has_vaild_right_node(&self) -> bool {
        self.right_ref().is_some()
    }
}

pub fn write_key_page(
    search_page: NodeWriteGuard,
    search_ref: &NodeCellRef,
    key: &EntryKey,
) -> (NodeWriteGuard, NodeCellRef) {
    // return search_page;
    if search_page.len() > 0 && search_page.last_key() < key {
        let right_ref = search_page.right_ref().unwrap();
        let right_node = right_ref.write();
        if !right_node.is_none() && right_node.len() > 0 && right_node.first_key() < key {
            debug!("will write {:?} to right page, start with {:?}", key, right_node.first_key());
            return write_key_page(right_node, right_ref, key);
        }
    }
    return (search_page, search_ref.clone());
}

const LATCH_FLAG: usize = !(!0 >> 1);

pub struct Node {
    data: UnsafeCell<NodeData>,
    cc: AtomicUsize,
}

impl Node {
    pub fn new(data: NodeData) -> Self {
        Node {
            data: UnsafeCell::new(data),
            cc: AtomicUsize::new(0),
        }
    }

    pub fn internal(innode: InNode) -> Self {
        Self::new(NodeData::Internal(box innode))
    }

    pub fn external(extnode: ExtNode) -> Self {
        Self::new(NodeData::External(box extnode))
    }

    pub fn none() -> Self {
        Self::new(NodeData::None)
    }
    pub fn none_ref() -> NodeCellRef {
        Arc::new(Node::none())
    }
    pub fn new_external(id: Id) -> Self {
        Self::external(ExtNode::new(id))
    }

    pub fn write(&self) -> NodeWriteGuard {
        debug!("acquiring node write lock");
        let cc = &self.cc;
        loop {
            let cc_num = cc.load(Relaxed);
            let expected = cc_num & (!LATCH_FLAG);
            debug_assert_eq!(expected & LATCH_FLAG, 0);
            match cc.compare_exchange_weak(expected, cc_num | LATCH_FLAG, Acquire, Relaxed) {
                Ok(num) if num == expected => {
                    return NodeWriteGuard {
                        data: self.data.get(),
                        cc: &self.cc as *const AtomicUsize,
                        version: node_version(cc_num),
                    }
                }
                _ => {}
            }
        }
    }

    pub fn read<'a, F: FnMut(&NodeReadHandler) -> R + 'a, R: 'a>(&'a self, mut func: F) -> R {
        let mut handler = NodeReadHandler {
            ptr: self.data.get(),
            version: 0,
        };
        let cc = &self.cc;
        loop {
            let cc_num = cc.load(SeqCst);
            if cc_num & LATCH_FLAG == LATCH_FLAG {
                debug!("read have a latch, retry {:b}", cc_num);
                continue;
            }
            handler.version = cc_num & (!LATCH_FLAG);
            let res = func(&handler);
            let new_cc_num = cc.load(SeqCst);
            if new_cc_num == cc_num {
                return res;
            }
            debug!("read version changed, retry {:b}", cc_num);
        }
    }

    pub fn read_unchecked(&self) -> &NodeData {
        unsafe { &*self.data.get() }
    }

    pub fn version(&self) -> usize {
        node_version(self.cc.load(SeqCst))
    }
}

pub fn node_version(cc_num: usize) -> usize {
    cc_num & (!LATCH_FLAG)
}

pub struct NodeWriteGuard {
    pub data: *mut NodeData,
    pub cc: *const AtomicUsize,
    pub version: usize,
}

impl Deref for NodeWriteGuard {
    type Target = NodeData;

    fn deref(&self) -> &<Self as Deref>::Target {
        debug_assert_ne!(self.data as usize, 0);
        unsafe { &*self.data }
    }
}

impl DerefMut for NodeWriteGuard {
    fn deref_mut(&mut self) -> &mut <Self as Deref>::Target {
        debug_assert_ne!(self.data as usize, 0);
        unsafe { &mut *self.data }
    }
}

impl Drop for NodeWriteGuard {
    fn drop(&mut self) {
        // cope with null pointer
        if self.cc as usize != 0 {
            let cc = unsafe { &*self.cc };
            let cc_num = cc.load(SeqCst);
            debug_assert_eq!(cc_num & LATCH_FLAG, LATCH_FLAG);
            cc.store((cc_num & (!LATCH_FLAG)) + 1, Release);
        }
    }
}

impl Default for NodeWriteGuard {
    fn default() -> Self {
        NodeWriteGuard {
            data: 0 as *mut NodeData,
            cc: 0 as *const AtomicUsize,
            version: 0,
        }
    }
}

unsafe impl Sync for Node {}

pub struct NodeReadHandler {
    pub ptr: *const NodeData,
    pub version: usize,
}

impl Deref for NodeReadHandler {
    type Target = NodeData;

    fn deref(&self) -> &<Self as Deref>::Target {
        unsafe { &*self.ptr }
    }
}

pub struct RemoveStatus {
    pub item_found: bool,
    pub removed: bool,
}

impl Default for Node {
    fn default() -> Self {
        Node::none()
    }
}

pub fn insert_into_split<T, S>(
    item: T,
    x: &mut S,
    y: &mut S,
    xlen: &mut usize,
    ylen: &mut usize,
    pos: usize,
) where
    S: Slice<Item = T> + BTreeSlice<T>,
    T: Default,
{
    debug!(
        "insert into split left len {}, right len {}, pos {}",
        xlen, ylen, pos
    );
    if pos < *xlen {
        debug!("insert into left part, pos: {}", pos);
        x.insert_at(item, pos, xlen);
    } else {
        let right_pos = pos - *xlen;
        debug!("insert into right part, pos: {}", right_pos);
        y.insert_at(item, right_pos, ylen);
    }
}
<commit_msg>simplify latch free right key search<commit_after>use super::*;
use std::sync::atomic::fence;
use std::sync::atomic::Ordering::AcqRel;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Release;

pub enum InsertSearchResult {
    External,
    Internal(NodeCellRef),
    RightNode(NodeCellRef),
}

pub enum InsertToNodeResult {
    NoSplit,
    Split(NodeWriteGuard, NodeCellRef, Option<EntryKey>),
    SplitParentChanged,
}

pub enum RemoveSearchResult {
    External,
    RightNode(NodeCellRef),
    Internal(NodeCellRef),
}

pub enum SubNodeStatus {
    ExtNodeEmpty,
    InNodeEmpty,
    Relocate(usize, usize),
    Merge(usize, usize),
    Ok,
}

pub struct NodeSplit {
    pub new_right_node: NodeCellRef,
    pub left_node_latch: NodeWriteGuard,
    pub pivot: EntryKey,
    pub parent_latch: NodeWriteGuard,
}

pub struct RebalancingNodes {
    pub left_guard: NodeWriteGuard,
    pub left_ref: NodeCellRef,
    pub right_guard: NodeWriteGuard,
    pub right_right_guard: NodeWriteGuard, // for external pointer modification
    pub parent: NodeWriteGuard,
    pub parent_pos: usize,
}

pub struct RemoveResult {
    pub rebalancing: Option<RebalancingNodes>,
    pub empty: bool,
    pub removed: bool,
}

pub enum NodeData {
    External(Box<ExtNode>),
    Internal(Box<InNode>),
    Empty(Box<EmptyNode>),
    None,
}

pub struct EmptyNode {
    pub left: Option<NodeCellRef>,
    pub right: NodeCellRef,
}

impl NodeData {
    pub fn is_none(&self) -> bool {
        match self {
            &NodeData::None => true,
            _ => false,
        }
    }
    pub fn is_empty(&self) -> bool {
        match self {
            &NodeData::Empty(ref n) => true,
            &NodeData::External(ref n) => n.len == 0,
            &NodeData::Internal(ref n) => n.len == 0,
            &NodeData::None => unreachable!(),
        }
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        let len = self.len();
        if self.is_ext() {
            self.extnode().search(key)
        } else {
            self.innode().search(key)
        }
    }

    pub fn remove(&mut self, pos: usize) {
        match self {
            &mut NodeData::External(ref mut node) => node.remove_at(pos),
            &mut NodeData::Internal(ref mut node) => node.remove_at(pos),
            &mut NodeData::None | &mut NodeData::Empty(_) => unreachable!(),
        }
    }
    pub fn is_ext(&self) -> bool {
        match self {
            &NodeData::External(_) => true,
            &NodeData::Internal(_) => false,
            &NodeData::None | &NodeData::Empty(_) => panic!(),
        }
    }
    pub fn keys(&self) -> &[EntryKey] {
        if self.is_ext() {
            &self.extnode().keys
        } else {
            &self.innode().keys
        }
    }

    pub fn first_key(&self) -> &EntryKey {
        &self.keys()[0]
    }

    pub fn last_key(&self) -> &EntryKey {
        &self.keys()[self.len() - 1]
    }

    pub fn len(&self) -> usize {
        match self {
            &NodeData::Internal(ref node) => node.len,
            &NodeData::External(ref node) => node.len,
            &NodeData::Empty(_) => 0,
            &NodeData::None => unreachable!(),
        }
    }

    // check if the node will be half full after an item have been removed
    pub fn is_half_full(&self) -> bool {
        if self.is_none() {
            true
        } else {
            let len = self.len();
            if len == 0 {
                return false;
            }
            len >= NUM_KEYS / 2 && len > 1
        }
    }

    pub fn cannot_merge(&self) -> bool {
        self.len() > NUM_KEYS / 2
    }

    pub fn extnode_mut(&mut self) -> &mut ExtNode {
        match self {
            &mut NodeData::External(ref mut node) => node,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn innode_mut(&mut self) -> &mut InNode {
        match self {
            &mut NodeData::Internal(ref mut n) => n,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn extnode(&self) -> &ExtNode {
        match self {
            &NodeData::External(ref node) => node,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn ext_id(&self) -> Id {
        match self {
            &NodeData::External(ref node) => node.id,
            &NodeData::None => Id::unit_id(),
            &NodeData::Internal(_) | &NodeData::Empty(_) => unreachable!(self.type_name()),
        }
    }
    pub fn innode(&self) -> &InNode {
        match self {
            &NodeData::Internal(ref n) => n,
            _ => unreachable!(self.type_name()),
        }
    }
    pub fn type_name(&self) -> &'static str {
        match self {
            &NodeData::Internal(_) => "internal",
            &NodeData::External(_) => "external",
            &NodeData::None => "none",
            &NodeData::Empty(_) => "empty",
        }
    }
    pub fn key_at_right_node(&self, key: &EntryKey) -> Option<&NodeCellRef> {
        if self.len() > 0 && self.last_key() < key {
            let right_node = self.right_ref().unwrap().read_unchecked();
            if !right_node.is_none() && right_node.len() > 0 && right_node.first_key() < key {
                debug!(
                    "found key to put to right page {:?}/{:?}",
                    key, right_node.first_key()
                );
                return Some(self.right_ref().unwrap())
            }
        }
        return None;
    }

    pub fn left_ref_mut(&mut self) -> Option<&mut NodeCellRef> {
        match self {
            &mut NodeData::External(ref mut n) => Some(&mut n.prev),
            _ => None,
        }
    }

    pub fn left_ref(&self) -> Option<&NodeCellRef> {
        match self {
            &NodeData::External(ref n) => Some(&n.prev),
            _ => None,
        }
    }

    pub fn right_ref_mut(&mut self) -> Option<&mut NodeCellRef> {
        match self {
            &mut NodeData::External(ref mut n) => Some(&mut n.next),
            &mut NodeData::Internal(ref mut n) => Some(&mut n.right),
            &mut NodeData::Empty(ref mut n) => Some(&mut n.right),
            &mut NodeData::None => None,
        }
    }

    pub fn right_ref(&self) -> Option<&NodeCellRef> {
        match self {
            &NodeData::External(ref n) => Some(&n.next),
            &NodeData::Internal(ref n) => Some(&n.right),
            &NodeData::Empty(ref n) => Some(&n.right),
            &NodeData::None => None,
        }
    }

    pub fn has_vaild_right_node(&self) -> bool {
        self.right_ref().is_some()
    }
}

pub fn write_key_page(
    search_page: NodeWriteGuard,
    search_ref: &NodeCellRef,
    key: &EntryKey,
) -> (NodeWriteGuard, NodeCellRef) {
    if search_page.len() > 0 && search_page.last_key() < key {
        let right_ref = search_page.right_ref().unwrap();
        let right_node = right_ref.write();
        if !right_node.is_none() && right_node.len() > 0 && right_node.first_key() < key {
            debug!("will write {:?} to right page, start with {:?}", key, right_node.first_key());
            return write_key_page(right_node, right_ref, key);
        }
    }
    return (search_page, search_ref.clone());
}

const LATCH_FLAG: usize = !(!0 >> 1);

pub struct Node {
    data: UnsafeCell<NodeData>,
    cc: AtomicUsize,
}

impl Node {
    pub fn new(data: NodeData) -> Self {
        Node {
            data: UnsafeCell::new(data),
            cc: AtomicUsize::new(0),
        }
    }

    pub fn internal(innode: InNode) -> Self {
        Self::new(NodeData::Internal(box innode))
    }

    pub fn external(extnode: ExtNode) -> Self {
        Self::new(NodeData::External(box extnode))
    }

    pub fn none() -> Self {
        Self::new(NodeData::None)
    }
    pub fn none_ref() -> NodeCellRef {
        Arc::new(Node::none())
    }
    pub fn new_external(id: Id) -> Self {
        Self::external(ExtNode::new(id))
    }

    pub fn write(&self) -> NodeWriteGuard {
        // debug!("acquiring node write lock");
        let cc = &self.cc;
        loop {
            let cc_num = cc.load(Relaxed);
            let expected = cc_num & (!LATCH_FLAG);
            debug_assert_eq!(expected & LATCH_FLAG, 0);
            match cc.compare_exchange_weak(expected, cc_num | LATCH_FLAG, Acquire, Relaxed) {
                Ok(num) if num == expected => {
                    return NodeWriteGuard {
                        data: self.data.get(),
                        cc: &self.cc as *const AtomicUsize,
                        version: node_version(cc_num),
                    }
                }
                _ => {}
            }
        }
    }

    pub fn read<'a, F: FnMut(&NodeReadHandler) -> R + 'a, R: 'a>(&'a self, mut func: F) -> R {
        let mut handler = NodeReadHandler {
            ptr: self.data.get(),
            version: 0,
        };
        let cc = &self.cc;
        loop {
            let cc_num = cc.load(SeqCst);
            if cc_num & LATCH_FLAG == LATCH_FLAG {
                // debug!("read have a latch, retry {:b}", cc_num);
                continue;
            }
            handler.version = cc_num & (!LATCH_FLAG);
            let res = func(&handler);
            let new_cc_num = cc.load(SeqCst);
            if new_cc_num == cc_num {
                return res;
            }
            // debug!("read version changed, retry {:b}", cc_num);
        }
    }

    pub fn read_unchecked(&self) -> &NodeData {
        unsafe { &*self.data.get() }
    }

    pub fn version(&self) -> usize {
        node_version(self.cc.load(SeqCst))
    }
}

pub fn node_version(cc_num: usize) -> usize {
    cc_num & (!LATCH_FLAG)
}

pub struct NodeWriteGuard {
    pub data: *mut NodeData,
    pub cc: *const AtomicUsize,
    pub version: usize,
}

impl Deref for NodeWriteGuard {
    type Target = NodeData;

    fn deref(&self) -> &<Self as Deref>::Target {
        debug_assert_ne!(self.data as usize, 0);
        unsafe { &*self.data }
    }
}

impl DerefMut for NodeWriteGuard {
    fn deref_mut(&mut self) -> &mut <Self as Deref>::Target {
        debug_assert_ne!(self.data as usize, 0);
        unsafe { &mut *self.data }
    }
}

impl Drop for NodeWriteGuard {
    fn drop(&mut self) {
        // cope with null pointer
        if self.cc as usize != 0 {
            let cc = unsafe { &*self.cc };
            let cc_num = cc.load(SeqCst);
            debug_assert_eq!(cc_num & LATCH_FLAG, LATCH_FLAG);
            cc.store((cc_num & (!LATCH_FLAG)) + 1, Release);
        }
    }
}

impl Default for NodeWriteGuard {
    fn default() -> Self {
        NodeWriteGuard {
            data: 0 as *mut NodeData,
            cc: 0 as *const AtomicUsize,
            version: 0,
        }
    }
}

unsafe impl Sync for Node {}

pub struct NodeReadHandler {
    pub ptr: *const NodeData,
    pub version: usize,
}

impl Deref for NodeReadHandler {
    type Target = NodeData;

    fn deref(&self) -> &<Self as Deref>::Target {
        unsafe { &*self.ptr }
    }
}

pub struct RemoveStatus {
    pub item_found: bool,
    pub removed: bool,
}

impl Default for Node {
    fn default() -> Self {
        Node::none()
    }
}

pub fn insert_into_split<T, S>(
    item: T,
    x: &mut S,
    y: &mut S,
    xlen: &mut usize,
    ylen: &mut usize,
    pos: usize,
) where
    S: Slice<Item = T> + BTreeSlice<T>,
    T: Default,
{
    debug!(
        "insert into split left len {}, right len {}, pos {}",
        xlen, ylen, pos
    );
    if pos < *xlen {
        debug!("insert into left part, pos: {}", pos);
        x.insert_at(item, pos, xlen);
    } else {
        let right_pos = pos - *xlen;
        debug!("insert into right part, pos: {}", right_pos);
        y.insert_at(item, right_pos, ylen);
    }
}
<|endoftext|>


<commit_before>use super::super::{Grid, Coord};
use super::*;

#[derive(Debug, Clone, Copy, PartialEq)]
enum EvCand {
    Ok,
    Elim(f64),
}

#[derive(Debug)]
enum Move {
    Decide(f64, usize, i32),
    Elim(f64, Vec<(usize, i32)>),
}

pub struct Evaluator {
    n_cells: usize,
    shape: FieldShape,
    val: Vec<i32>,
    clue: Vec<i32>,
    cand_score: Vec<[EvCand; (MAX_VAL + 1) as usize]>, // 1-origin for simplicity
    total_score: f64,
    inconsistent: bool,
    move_cand: Vec<Move>,
}

const UNIQUE_ELIMINATION: f64 = 0.8f64;
const SMALL_LARGE_ELIMINATION: f64 = 1.0f64;

impl Evaluator {
    pub fn new(problem: &Grid<Clue>) -> Evaluator {
        let n_cells = (problem.height() * problem.width()) as usize;

        let mut has_clue = Grid::new(problem.height(), problem.width(), false);
        for i in 0..n_cells {
            if let Clue::Clue { vertical: _, horizontal: _ } = problem[i] {
                has_clue[i] = true;
            }
        }
        let shape = FieldShape::new(&has_clue);
        let mut clue = vec![0; shape.group_to_cells.len()];
        for i in 0..shape.group_to_cells.len() {
            let clue_loc = shape.clue_locations[i];
            clue[i] = match clue_loc {
                ClueLocation::Horizontal(h) => match problem[h] {
                    Clue::Clue { horizontal: h, vertical: _ } => h,
                    _ => unreachable!(),
                },
                ClueLocation::Vertical(v) => match problem[v] {
                    Clue::Clue { horizontal: _, vertical: v } => v,
                    _ => unreachable!(),
                },
            }
        }
        Evaluator {
            n_cells: n_cells,
            shape: shape,
            val: vec![UNDECIDED; n_cells],
            clue: clue,
            cand_score: vec![[EvCand::Ok; (MAX_VAL + 1) as usize]; n_cells],
            total_score: 0.0f64,
            inconsistent: false,
            move_cand: vec![],
        }
    }
    pub fn evaluate(&mut self) -> Option<f64> {
        let mut n_undecided = 0;
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] {
                n_undecided += 1;
            }
        }
        loop {
            self.move_cand.clear();

            // apply possible eliminations
            self.simple_elimination();
            self.unique_decision();
            self.unique_decision();
            self.two_cells_propagation();
            self.naked_pair();

            self.remove_unnecessary_move();

            if self.move_cand.len() == 0 {
                break;
            }

            let mut lowest_decision = (1e20f64, 0, 0);
            let mut lowest_elim = (1e20f64, vec![]);
            for mv in &self.move_cand {
                match mv {
                    &Move::Decide(score, loc, v) => if score < lowest_decision.0 {
                        lowest_decision = (score, loc, v);
                    },
                    &Move::Elim(score, ref elims) => if score < lowest_elim.0 {
                        lowest_elim = (score, elims.clone());
                    }
                }
            }
            if lowest_elim.0 < lowest_decision.0 * 11.0f64 {
                // use elim
                for (loc, v) in lowest_elim.1 {
                    self.cand_score[loc][v as usize] = EvCand::Elim(lowest_elim.0);
                }
            } else {
                // use lowest_decision
                let mut move_score = 0.0f64;
                for mv in &self.move_cand {
                    if let &Move::Decide(score, _, _) = mv {
                        move_score += score.powf(-2.0f64);
                    }
                }
                move_score = move_score.powf(-1.0 / 2.0f64);
                move_score *= (n_undecided as f64).powf(0.3);

                self.total_score += move_score;
                self.val[lowest_decision.1] = lowest_decision.2;
                n_undecided -= 1;
            }
        }
        if n_undecided == 0 {
            Some(self.total_score)
        } else {
            None
        }
    }
    fn remove_unnecessary_move(&mut self) {
        let mut new_cand = vec![];
        for m in &self.move_cand {
            match m {
                &Move::Decide(score, loc, v) =>
                    if self.val[loc] == UNDECIDED {
                        new_cand.push(Move::Decide(score, loc, v))
                    },
                &Move::Elim(score, ref elims) => {
                    let mut new_elims = vec![];
                    for &(loc, v) in elims {
                        let isok = match self.cand_score[loc][v as usize] {
                            EvCand::Ok => true,
                            EvCand::Elim(sc) => score < sc,
                        };
                        if isok {
                            new_elims.push((loc, v));
                        }
                    }
                    if new_elims.len() > 0 {
                        new_cand.push(Move::Elim(score, new_elims));
                    }
                },
            }
        }
        self.move_cand = new_cand;
    }
    fn simple_elimination(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut used_mask = 0u32; // bit indices are 1-origin
            let mut rem_sum = self.clue[gi];
            for c in self.shape.group_to_cells[gi] {
                if self.val[c] != UNDECIDED {
                    used_mask |= 1u32 << self.val[c];
                    rem_sum -= self.val[c];
                }
            }
            let rem_cells = self.shape.group_to_cells[gi].size() as i32 - used_mask.count_ones() as i32;
            if rem_cells == 0 {
                continue;
            }
            if rem_cells == 1 {
                let mut ep = 0;
                for c in self.shape.group_to_cells[gi] {
                    if self.val[c] == UNDECIDED {
                        ep = c;
                    }
                }
                self.move_cand.push(Move::Decide(4.0, ep, rem_sum));
                continue;
            }

            let mut allowed_cand = ((2u32 << MAX_VAL) - 2u32) ^ used_mask;
            let mut required = 0u32;
            // eliminate too large candidates
            {
                let mut sum_small = 0;
                let mut n_cand_checked = 0;
                let mut max_cand_checked = 0;
                let mut low_bits = 0;
                for i in 1..(MAX_VAL + 1) {
                    if n_cand_checked == rem_cells - 1 {
                        break;
                    }
                    if (allowed_cand & (1u32 << i)) != 0 {
                        sum_small += i;
                        n_cand_checked += 1;
                        max_cand_checked = i;
                        low_bits |= 1u32 << i;
                    }
                }
                let mut max_allowed = rem_sum - sum_small;

                if max_allowed <= max_cand_checked {
                    self.inconsistent = true;
                    return;
                }
                if max_allowed < MAX_VAL {
                    if max_allowed == max_cand_checked + 2 {
                        allowed_cand &= !(1u32 << (max_cand_checked + 1));
                    }
                    allowed_cand &= (2u32 << max_allowed) - 2u32;
                }
                if max_allowed == max_cand_checked + 1 {
                    required |= low_bits | (1u32 << (max_cand_checked + 1));
                } else if max_allowed == max_cand_checked + 2 {
                    required |= low_bits | (1u32 << (max_cand_checked + 2));
                }
            }
            {
                let mut sum_large = 0;
                let mut n_cand_checked = 0;
                let mut min_cand_checked = 0;
                let mut high_bits = 0u32;
                for i in 1..(MAX_VAL + 1) {
                    let i = MAX_VAL + 1 - i;
                    if n_cand_checked == rem_cells - 1 {
                        break;
                    }
                    if (allowed_cand & (1u32 << i)) != 0 {
                        sum_large += i;
                        n_cand_checked += 1;
                        min_cand_checked = i;
                        high_bits |= 1u32 << i;
                    }
                }
                let mut min_allowed = rem_sum - sum_large;

                if min_allowed >= min_cand_checked {
                    self.inconsistent = true;
                    return;
                }
                if min_allowed > 1 {
                    if min_allowed == min_cand_checked - 2 {
                        allowed_cand &= !(1u32 << (min_cand_checked - 1));
                    }
                    allowed_cand &= !((2u32 << (min_allowed - 1)) - 2u32);
                }
                if min_allowed == min_cand_checked - 1 {
                    required |= high_bits | (1u32 << (min_cand_checked - 1));
                } else if min_allowed == min_cand_checked - 2 {
                    required |= high_bits | (1u32 << (min_cand_checked - 2));
                }
            }

            let mut elims = vec![];
            for n in 1..(MAX_VAL + 1) {
                if (allowed_cand & (1u32 << n)) == 0 {
                    for c in self.shape.group_to_cells[gi] {
                        if self.val[c] == UNDECIDED {
                            elims.push((c, n));
                        }
                    }
                }
            }
            self.move_cand.push(Move::Elim(SMALL_LARGE_ELIMINATION, elims));

            for n in 1..(MAX_VAL + 1) {
                if (required & (1u32 << n)) != 0 {
                    let mut cost = 0.0f64;
                    let mut pos = None;
                    let mut twice = false;
                    for c in self.shape.group_to_cells[gi] {
                        if self.val[c] == UNDECIDED {
                            if let EvCand::Elim(s) = self.cand_score[c][n as usize] {
                                cost += s;
                            } else {
                                if pos == None {
                                    pos = Some(c);
                                } else {
                                    twice = true;
                                    break;
                                }
                            }
                        } else if self.val[c] == n {
                            twice = true;
                            break;
                        }
                    }
                    if !twice {
                        if let Some(c) = pos {
                            self.move_cand.push(Move::Decide(cost + 4.0, c, n));
                        } else {
                            self.inconsistent = true;
                            return;
                        }
                    }
                }
            }
        }
    }
    fn unique_elimination(&mut self) {
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] && self.val[i] != UNDECIDED {
                let (g1, g2) = self.shape.cell_to_groups[i];
                let mut mv = vec![];
                for c in self.shape.group_to_cells[g1] {
                    if c != i {
                        mv.push((c, self.val[i]));
                    }
                }
                for c in self.shape.group_to_cells[g2] {
                    if c != i {
                        mv.push((c, self.val[i]));
                    }
                }
                self.move_cand.push(Move::Elim(UNIQUE_ELIMINATION, mv));
            }
        }
    }
    fn unique_decision(&mut self) {
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] {
                let mut cost = 0.0f64;
                let mut cand = -1;
                for v in 1..(MAX_VAL + 1) {
                    match self.cand_score[i][v as usize] {
                        EvCand::Ok => if cand == -1 {
                            cand = v;
                        } else {
                            cand = -2;
                        },
                        EvCand::Elim(c) => cost += c,
                    }
                }
                if cand == -1 {
                    self.inconsistent = true;
                    return;
                } else if cand != -2 {
                    self.move_cand.push(Move::Decide(cost, i, cand));
                }
            }
        }
    }
    fn two_cells_propagation(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut used_mask = 0u32; // bit indices are 1-origin
            let mut rem_sum = self.clue[gi];
            let mut undet1 = None;
            let mut undet2 = None;
            let mut more_than_two = false;
            for c in self.shape.group_to_cells[gi] {
                if self.val[c] == UNDECIDED {
                    if undet1 == None {
                        undet1 = Some(c);
                    } else if undet2 == None {
                        undet2 = Some(c);
                    } else {
                        more_than_two = true;
                    }
                } else {
                    rem_sum -= self.val[c];
                }
            }
            if !more_than_two && undet2 != None {
                if let (Some(c1), Some(c2)) = (undet1, undet2) {
                    for n in 1..(MAX_VAL + 1) {
                        if let EvCand::Elim(s) = self.cand_score[c1][n as usize] {
                            let n2 = rem_sum - n;
                            if 1 <= n2 && n2 <= MAX_VAL {
                                self.move_cand.push(Move::Elim(s + 5.0f64, vec![(c2, n2)]));
                            }
                        }
                        if let EvCand::Elim(s) = self.cand_score[c2][n as usize] {
                            let n2 = rem_sum - n;
                            if 1 <= n2 && n2 <= MAX_VAL {
                                self.move_cand.push(Move::Elim(s + 5.0f64, vec![(c1, n2)]));
                            }
                        }
                    }
                } else {
                    unreachable!();
                }
            }
        }
    }
    fn naked_pair(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut two_cand_cells = vec![];
            for c1 in self.shape.group_to_cells[gi] {
                let mut oks = 0u32;
                let mut score = 0.0f64;
                for n in 1..(MAX_VAL + 1) {
                    if let EvCand::Elim(s) = self.cand_score[c1][n as usize] {
                        score += s;
                    } else {
                        oks |= 1u32 << n;
                    }
                }
                if oks.count_ones() == 2 {
                    two_cand_cells.push((c1, oks, score));
                }
            }
            for i in 0..two_cand_cells.len() {
                for j in (i + 1)..two_cand_cells.len() {
                    if two_cand_cells[i].1 == two_cand_cells[j].1 {
                        let mut elims = vec![];
                        let x = two_cand_cells[i].1.trailing_zeros() as i32;
                        let y = (two_cand_cells[i].1 ^ (1u32 << x)).trailing_zeros() as i32;
                        for c in self.shape.group_to_cells[gi] {
                            if two_cand_cells[i].0 != c && two_cand_cells[j].0 != c {
                                elims.push((c, x));
                                elims.push((c, y));
                            }
                        }
                        self.move_cand.push(Move::Elim(two_cand_cells[i].2 + two_cand_cells[j].2, elims));
                    }
                }
            }
        }
    }
}
<commit_msg>Kakuro: enhance ability of slimple_elimination in Evaluator<commit_after>use super::super::{Grid, Coord};
use super::*;

#[derive(Debug, Clone, Copy, PartialEq)]
enum EvCand {
    Ok,
    Elim(f64),
}

#[derive(Debug)]
enum Move {
    Decide(f64, usize, i32),
    Elim(f64, Vec<(usize, i32)>),
}

pub struct Evaluator {
    n_cells: usize,
    shape: FieldShape,
    val: Vec<i32>,
    clue: Vec<i32>,
    cand_score: Vec<[EvCand; (MAX_VAL + 1) as usize]>, // 1-origin for simplicity
    total_score: f64,
    inconsistent: bool,
    move_cand: Vec<Move>,
}

const UNIQUE_ELIMINATION: f64 = 0.8f64;
const SMALL_LARGE_ELIMINATION: f64 = 1.0f64;

impl Evaluator {
    pub fn new(problem: &Grid<Clue>) -> Evaluator {
        let n_cells = (problem.height() * problem.width()) as usize;

        let mut has_clue = Grid::new(problem.height(), problem.width(), false);
        for i in 0..n_cells {
            if let Clue::Clue { vertical: _, horizontal: _ } = problem[i] {
                has_clue[i] = true;
            }
        }
        let shape = FieldShape::new(&has_clue);
        let mut clue = vec![0; shape.group_to_cells.len()];
        for i in 0..shape.group_to_cells.len() {
            let clue_loc = shape.clue_locations[i];
            clue[i] = match clue_loc {
                ClueLocation::Horizontal(h) => match problem[h] {
                    Clue::Clue { horizontal: h, vertical: _ } => h,
                    _ => unreachable!(),
                },
                ClueLocation::Vertical(v) => match problem[v] {
                    Clue::Clue { horizontal: _, vertical: v } => v,
                    _ => unreachable!(),
                },
            }
        }
        Evaluator {
            n_cells: n_cells,
            shape: shape,
            val: vec![UNDECIDED; n_cells],
            clue: clue,
            cand_score: vec![[EvCand::Ok; (MAX_VAL + 1) as usize]; n_cells],
            total_score: 0.0f64,
            inconsistent: false,
            move_cand: vec![],
        }
    }
    pub fn evaluate(&mut self) -> Option<f64> {
        let mut n_undecided = 0;
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] {
                n_undecided += 1;
            }
        }
        loop {
            self.move_cand.clear();

            // apply possible eliminations
            self.simple_elimination();
            self.unique_decision();
            self.unique_decision();
            self.two_cells_propagation();
            self.naked_pair();

            self.remove_unnecessary_move();

            if self.move_cand.len() == 0 {
                break;
            }

            let mut lowest_decision = (1e20f64, 0, 0);
            let mut lowest_elim = (1e20f64, vec![]);
            for mv in &self.move_cand {
                match mv {
                    &Move::Decide(score, loc, v) => if score < lowest_decision.0 {
                        lowest_decision = (score, loc, v);
                    },
                    &Move::Elim(score, ref elims) => if score < lowest_elim.0 {
                        lowest_elim = (score, elims.clone());
                    }
                }
            }
            if lowest_elim.0 < lowest_decision.0 * 11.0f64 {
                // use elim
                for (loc, v) in lowest_elim.1 {
                    self.cand_score[loc][v as usize] = EvCand::Elim(lowest_elim.0);
                }
            } else {
                // use lowest_decision
                let mut move_score = 0.0f64;
                for mv in &self.move_cand {
                    if let &Move::Decide(score, _, _) = mv {
                        move_score += score.powf(-2.0f64);
                    }
                }
                move_score = move_score.powf(-1.0 / 2.0f64);
                move_score *= (n_undecided as f64).powf(0.3);

                self.total_score += move_score;
                self.val[lowest_decision.1] = lowest_decision.2;
                n_undecided -= 1;
            }
        }
        if n_undecided == 0 {
            Some(self.total_score)
        } else {
            None
        }
    }
    fn remove_unnecessary_move(&mut self) {
        let mut new_cand = vec![];
        for m in &self.move_cand {
            match m {
                &Move::Decide(score, loc, v) =>
                    if self.val[loc] == UNDECIDED {
                        new_cand.push(Move::Decide(score, loc, v))
                    },
                &Move::Elim(score, ref elims) => {
                    let mut new_elims = vec![];
                    for &(loc, v) in elims {
                        let isok = match self.cand_score[loc][v as usize] {
                            EvCand::Ok => true,
                            EvCand::Elim(sc) => score < sc,
                        };
                        if isok {
                            new_elims.push((loc, v));
                        }
                    }
                    if new_elims.len() > 0 {
                        new_cand.push(Move::Elim(score, new_elims));
                    }
                },
            }
        }
        self.move_cand = new_cand;
    }
    fn simple_elimination(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut used_mask = 0u32; // bit indices are 1-origin
            let mut rem_sum = self.clue[gi];
            for c in self.shape.group_to_cells[gi] {
                if self.val[c] != UNDECIDED {
                    used_mask |= 1u32 << self.val[c];
                    rem_sum -= self.val[c];
                }
            }
            let rem_cells = self.shape.group_to_cells[gi].size() as i32 - used_mask.count_ones() as i32;
            if rem_cells == 0 {
                continue;
            }
            if rem_cells == 1 {
                let mut ep = 0;
                for c in self.shape.group_to_cells[gi] {
                    if self.val[c] == UNDECIDED {
                        ep = c;
                    }
                }
                self.move_cand.push(Move::Decide(4.0, ep, rem_sum));
                continue;
            }

            let mut allowed_cand = ((2u32 << MAX_VAL) - 2u32) ^ used_mask;
            let mut required = 0u32;
            // eliminate too large candidates
            {
                let mut sum_small = 0;
                let mut n_cand_checked = 0;
                let mut max_cand_checked = 0;
                let mut low_bits = 0;
                for i in 1..(MAX_VAL + 1) {
                    if (allowed_cand & (1u32 << i)) != 0 {
                        if n_cand_checked == rem_cells - 1 {
                            max_cand_checked = i;
                            break;
                        } else {
                            sum_small += i;
                            n_cand_checked += 1;
                            low_bits |= 1u32 << i;
                        }
                    }
                }
                let mut max_allowed = rem_sum - sum_small;

                if max_allowed < max_cand_checked {
                    self.inconsistent = true;
                    return;
                }
                if max_allowed < MAX_VAL {
                    if max_allowed == max_cand_checked + 1 {
                        allowed_cand &= !(1u32 << max_cand_checked);
                    }
                    allowed_cand &= (2u32 << max_allowed) - 2u32;
                }
                if max_allowed == max_cand_checked {
                    required |= low_bits | (1u32 << max_cand_checked);
                } else if max_allowed == max_cand_checked {
                    required |= low_bits | (1u32 << (max_cand_checked + 1));
                }
            }
            {
                let mut sum_large = 0;
                let mut n_cand_checked = 0;
                let mut min_cand_checked = 0;
                let mut high_bits = 0u32;
                for i in 1..(MAX_VAL + 1) {
                    let i = MAX_VAL + 1 - i;
                    if (allowed_cand & (1u32 << i)) != 0 {
                        if n_cand_checked == rem_cells - 1 {
                            min_cand_checked = i;
                            break;
                        } else {
                            sum_large += i;
                            n_cand_checked += 1;
                            high_bits |= 1u32 << i;
                        }
                    }
                }
                let mut min_allowed = rem_sum - sum_large;

                if min_allowed > min_cand_checked {
                    self.inconsistent = true;
                    return;
                }
                if min_allowed > 1 {
                    if min_allowed == min_cand_checked - 1 {
                        allowed_cand &= !(1u32 << min_cand_checked);
                    }
                    allowed_cand &= !((2u32 << (min_allowed - 1)) - 2u32);
                }
                if min_allowed == min_cand_checked {
                    required |= high_bits | (1u32 << min_cand_checked);
                } else if min_allowed == min_cand_checked - 1 {
                    required |= high_bits | (1u32 << (min_cand_checked - 1));
                }
            }

            let mut elims = vec![];
            for n in 1..(MAX_VAL + 1) {
                if (allowed_cand & (1u32 << n)) == 0 {
                    for c in self.shape.group_to_cells[gi] {
                        if self.val[c] == UNDECIDED {
                            elims.push((c, n));
                        }
                    }
                }
            }
            self.move_cand.push(Move::Elim(SMALL_LARGE_ELIMINATION, elims));

            for n in 1..(MAX_VAL + 1) {
                if (required & (1u32 << n)) != 0 {
                    let mut cost = 0.0f64;
                    let mut pos = None;
                    let mut twice = false;
                    for c in self.shape.group_to_cells[gi] {
                        if self.val[c] == UNDECIDED {
                            if let EvCand::Elim(s) = self.cand_score[c][n as usize] {
                                cost += s;
                            } else {
                                if pos == None {
                                    pos = Some(c);
                                } else {
                                    twice = true;
                                    break;
                                }
                            }
                        } else if self.val[c] == n {
                            twice = true;
                            break;
                        }
                    }
                    if !twice {
                        if let Some(c) = pos {
                            self.move_cand.push(Move::Decide(cost + 4.0, c, n));
                        } else {
                            self.inconsistent = true;
                            return;
                        }
                    }
                }
            }
        }
    }
    fn unique_elimination(&mut self) {
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] && self.val[i] != UNDECIDED {
                let (g1, g2) = self.shape.cell_to_groups[i];
                let mut mv = vec![];
                for c in self.shape.group_to_cells[g1] {
                    if c != i {
                        mv.push((c, self.val[i]));
                    }
                }
                for c in self.shape.group_to_cells[g2] {
                    if c != i {
                        mv.push((c, self.val[i]));
                    }
                }
                self.move_cand.push(Move::Elim(UNIQUE_ELIMINATION, mv));
            }
        }
    }
    fn unique_decision(&mut self) {
        for i in 0..self.n_cells {
            if !self.shape.has_clue[i] {
                let mut cost = 0.0f64;
                let mut cand = -1;
                for v in 1..(MAX_VAL + 1) {
                    match self.cand_score[i][v as usize] {
                        EvCand::Ok => if cand == -1 {
                            cand = v;
                        } else {
                            cand = -2;
                        },
                        EvCand::Elim(c) => cost += c,
                    }
                }
                if cand == -1 {
                    self.inconsistent = true;
                    return;
                } else if cand != -2 {
                    self.move_cand.push(Move::Decide(cost, i, cand));
                }
            }
        }
    }
    fn two_cells_propagation(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut used_mask = 0u32; // bit indices are 1-origin
            let mut rem_sum = self.clue[gi];
            let mut undet1 = None;
            let mut undet2 = None;
            let mut more_than_two = false;
            for c in self.shape.group_to_cells[gi] {
                if self.val[c] == UNDECIDED {
                    if undet1 == None {
                        undet1 = Some(c);
                    } else if undet2 == None {
                        undet2 = Some(c);
                    } else {
                        more_than_two = true;
                    }
                } else {
                    rem_sum -= self.val[c];
                }
            }
            if !more_than_two && undet2 != None {
                if let (Some(c1), Some(c2)) = (undet1, undet2) {
                    for n in 1..(MAX_VAL + 1) {
                        if let EvCand::Elim(s) = self.cand_score[c1][n as usize] {
                            let n2 = rem_sum - n;
                            if 1 <= n2 && n2 <= MAX_VAL {
                                self.move_cand.push(Move::Elim(s + 5.0f64, vec![(c2, n2)]));
                            }
                        }
                        if let EvCand::Elim(s) = self.cand_score[c2][n as usize] {
                            let n2 = rem_sum - n;
                            if 1 <= n2 && n2 <= MAX_VAL {
                                self.move_cand.push(Move::Elim(s + 5.0f64, vec![(c1, n2)]));
                            }
                        }
                    }
                } else {
                    unreachable!();
                }
            }
        }
    }
    fn naked_pair(&mut self) {
        for gi in 0..self.shape.group_to_cells.len() {
            let mut two_cand_cells = vec![];
            for c1 in self.shape.group_to_cells[gi] {
                let mut oks = 0u32;
                let mut score = 0.0f64;
                for n in 1..(MAX_VAL + 1) {
                    if let EvCand::Elim(s) = self.cand_score[c1][n as usize] {
                        score += s;
                    } else {
                        oks |= 1u32 << n;
                    }
                }
                if oks.count_ones() == 2 {
                    two_cand_cells.push((c1, oks, score));
                }
            }
            for i in 0..two_cand_cells.len() {
                for j in (i + 1)..two_cand_cells.len() {
                    if two_cand_cells[i].1 == two_cand_cells[j].1 {
                        let mut elims = vec![];
                        let x = two_cand_cells[i].1.trailing_zeros() as i32;
                        let y = (two_cand_cells[i].1 ^ (1u32 << x)).trailing_zeros() as i32;
                        for c in self.shape.group_to_cells[gi] {
                            if two_cand_cells[i].0 != c && two_cand_cells[j].0 != c {
                                elims.push((c, x));
                                elims.push((c, y));
                            }
                        }
                        self.move_cand.push(Move::Elim(two_cand_cells[i].2 + two_cand_cells[j].2, elims));
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate bufstream;
extern crate cargo;
extern crate filetime;
extern crate flate2;
extern crate git2;
extern crate hamcrest;
extern crate libc;
extern crate rustc_serialize;
extern crate tar;
extern crate tempdir;
extern crate term;
extern crate url;
#[cfg(windows)] extern crate kernel32;
#[cfg(windows)] extern crate winapi;

#[macro_use]
extern crate log;

use cargo::util::Rustc;
use std::ffi::OsStr;

mod support;
macro_rules! test {
    ($name:ident $expr:expr) => (
        #[test]
        fn $name() {
            ::support::paths::setup();
            setup();
            $expr;
        }
    )
}

mod test_bad_config;
mod test_bad_manifest_path;
mod test_cargo;
mod test_cargo_bench;
mod test_cargo_build_auth;
mod test_cargo_build_lib;
mod test_cargo_clean;
mod test_cargo_compile;
mod test_cargo_compile_custom_build;
mod test_cargo_compile_git_deps;
mod test_cargo_compile_path_deps;
mod test_cargo_compile_plugins;
mod test_cargo_compile_rustflags;
mod test_cargo_cross_compile;
mod test_cargo_doc;
mod test_cargo_features;
mod test_cargo_fetch;
mod test_cargo_freshness;
mod test_cargo_generate_lockfile;
mod test_cargo_init;
mod test_cargo_install;
mod test_cargo_metadata;
mod test_cargo_new;
mod test_cargo_package;
mod test_cargo_profiles;
mod test_cargo_overrides;
mod test_cargo_publish;
mod test_cargo_read_manifest;
mod test_cargo_registry;
mod test_cargo_run;
mod test_cargo_concurrent;
mod test_cargo_rustc;
mod test_cargo_rustdoc;
mod test_cargo_search;
mod test_cargo_test;
mod test_cargo_tool_paths;
mod test_cargo_config;
mod test_cargo_verify_project;
mod test_cargo_version;
mod test_shell;
mod test_cargo_death;
mod test_cargo_cfg;

thread_local!(static RUSTC: Rustc = Rustc::new("rustc").unwrap());

fn rustc_host() -> String {
    RUSTC.with(|r| r.host.clone())
}

fn is_nightly() -> bool {
    RUSTC.with(|r| {
        r.verbose_version.contains("-nightly") ||
            r.verbose_version.contains("-dev")
    })
}

fn process<T: AsRef<OsStr>>(t: T) -> cargo::util::ProcessBuilder {
    let mut p = cargo::util::process(t.as_ref());
    p.cwd(&support::paths::root())
     .env("HOME", &support::paths::home())
     .env_remove("CARGO_HOME")
     .env_remove("RUSTC")
     .env_remove("XDG_CONFIG_HOME")      // see #2345
     .env("GIT_CONFIG_NOSYSTEM", "1")    // keep trying to sandbox ourselves
     .env_remove("CARGO_TARGET_DIR")     // we assume 'target'
     .env_remove("MSYSTEM");             // assume cmd.exe everywhere on windows
    return p
}

fn cargo_process() -> cargo::util::ProcessBuilder {
    process(&support::cargo_dir().join("cargo"))
}

#[allow(deprecated)] // sleep_ms is now deprecated in favor of sleep()
fn sleep_ms(ms: u32) {
    std::thread::sleep_ms(ms);
}
<commit_msg>Auto merge of #2617 - sanxiyn:reset-rustflags, r=alexcrichton<commit_after>#![deny(warnings)]

extern crate bufstream;
extern crate cargo;
extern crate filetime;
extern crate flate2;
extern crate git2;
extern crate hamcrest;
extern crate libc;
extern crate rustc_serialize;
extern crate tar;
extern crate tempdir;
extern crate term;
extern crate url;
#[cfg(windows)] extern crate kernel32;
#[cfg(windows)] extern crate winapi;

#[macro_use]
extern crate log;

use cargo::util::Rustc;
use std::ffi::OsStr;

mod support;
macro_rules! test {
    ($name:ident $expr:expr) => (
        #[test]
        fn $name() {
            ::support::paths::setup();
            setup();
            $expr;
        }
    )
}

mod test_bad_config;
mod test_bad_manifest_path;
mod test_cargo;
mod test_cargo_bench;
mod test_cargo_build_auth;
mod test_cargo_build_lib;
mod test_cargo_clean;
mod test_cargo_compile;
mod test_cargo_compile_custom_build;
mod test_cargo_compile_git_deps;
mod test_cargo_compile_path_deps;
mod test_cargo_compile_plugins;
mod test_cargo_compile_rustflags;
mod test_cargo_cross_compile;
mod test_cargo_doc;
mod test_cargo_features;
mod test_cargo_fetch;
mod test_cargo_freshness;
mod test_cargo_generate_lockfile;
mod test_cargo_init;
mod test_cargo_install;
mod test_cargo_metadata;
mod test_cargo_new;
mod test_cargo_package;
mod test_cargo_profiles;
mod test_cargo_overrides;
mod test_cargo_publish;
mod test_cargo_read_manifest;
mod test_cargo_registry;
mod test_cargo_run;
mod test_cargo_concurrent;
mod test_cargo_rustc;
mod test_cargo_rustdoc;
mod test_cargo_search;
mod test_cargo_test;
mod test_cargo_tool_paths;
mod test_cargo_config;
mod test_cargo_verify_project;
mod test_cargo_version;
mod test_shell;
mod test_cargo_death;
mod test_cargo_cfg;

thread_local!(static RUSTC: Rustc = Rustc::new("rustc").unwrap());

fn rustc_host() -> String {
    RUSTC.with(|r| r.host.clone())
}

fn is_nightly() -> bool {
    RUSTC.with(|r| {
        r.verbose_version.contains("-nightly") ||
            r.verbose_version.contains("-dev")
    })
}

fn process<T: AsRef<OsStr>>(t: T) -> cargo::util::ProcessBuilder {
    let mut p = cargo::util::process(t.as_ref());
    p.cwd(&support::paths::root())
     .env("HOME", &support::paths::home())
     .env_remove("CARGO_HOME")
     .env_remove("RUSTC")
     .env_remove("RUSTFLAGS")
     .env_remove("XDG_CONFIG_HOME")      // see #2345
     .env("GIT_CONFIG_NOSYSTEM", "1")    // keep trying to sandbox ourselves
     .env_remove("CARGO_TARGET_DIR")     // we assume 'target'
     .env_remove("MSYSTEM");             // assume cmd.exe everywhere on windows
    return p
}

fn cargo_process() -> cargo::util::ProcessBuilder {
    process(&support::cargo_dir().join("cargo"))
}

#[allow(deprecated)] // sleep_ms is now deprecated in favor of sleep()
fn sleep_ms(ms: u32) {
    std::thread::sleep_ms(ms);
}
<|endoftext|>


<commit_before>// Copyright 2018 Kieran Hervold
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Create a weight matrix representing a set of aligned reference sequences
//! that constitute a motif, and use this matrix to scan query sequences for
//! occurrences of this motif.
//! Complexity: O(n*m) for motif length n and query length m
//!
//! The position-specific scoring matrix (PSSM), aka position weight matrix (PWM),
//! algorithm is implemented for both DNA and amino-acid sequences.
//!
//! # Examples
//!
//! use bio::pattern_matching::pssm::DNAMotif;
//! let pssm = DNAMotif::from_seqs(vec![
//!            b"AAAA".to_vec(),
//!            b"AATA".to_vec(),
//!            b"AAGA".to_vec(),
//!            b"AAAA".to_vec(),
//!        ].as_ref(), None).unwrap();
//! let start_pos = pssm.score(b"CCCCCAATA").unwrap().loc;
//! println!("motif found at position {}", start_pos);
//!
//! /* amino acid sequences are supported, too */
//! use pssm::pattern_matching::pssm::ProtMotif;
//! let pssm = ProtMotif::from_seqs(vec![
//!            b"ARNNYM".to_vec(),
//!            b"ARNRYM".to_vec(),
//!            b"ARNNCM".to_vec(),
//!            b"ARNNYM".to_vec(),
//!        ].as_ref(), None).unwrap();

use std::borrow::Borrow;
use std::f32::NEG_INFINITY;

use itertools::Itertools;
use ndarray::prelude::Array2;

mod dnamotif;
pub mod errors;
mod protmotif;

pub use self::dnamotif::DNAMotif;
pub use self::errors::{Error, Result};
pub use self::protmotif::ProtMotif;

/// default pseudocount - used to prevent 0 tallies
pub const DEF_PSEUDO: f32 = 0.5;
/// approximately zero
pub const EPSILON: f32 = 1e-5;
/// value representing an invalid monomer in lookup table
pub const INVALID_MONO: u8 = 255;

/// Represents motif score & location of match
#[derive(Debug, Clone, PartialEq)]
pub struct ScoredPos {
    pub loc: usize,
    pub sum: f32,
    pub scores: Vec<f32>,
}

impl Default for ScoredPos {
    fn default() -> ScoredPos {
        ScoredPos {
            loc: 0,
            sum: NEG_INFINITY,
            scores: Vec::new(),
        }
    }
}

/// Trait containing code shared between DNA and protein implementations
/// of the position-specific scoring matrix.
pub trait Motif {
    /// Lookup table mapping monomer -> index
    const LK: [u8; 127] = [INVALID_MONO; 127];
    /// All monomers, in order corresponding to lookup table
    const MONOS: &'static [u8] = b"";
    /// Monomer count - equal to length of `MONOS`
    const MONO_CT: usize = 0;

    /// Returns a weight matrix representing the sequences provided.
    /// This code is shared by implementations of `from_seqs`
    /// # Arguments
    /// * `seqs` - sequences incorporated into motif
    /// * `pseudos` - array slice with a pseudocount for each monomer;
    ///    defaults to DEF_PSEUDO for all if None is supplied
    ///
    /// FIXME: pseudos should be an array of size MONO_CT, but that
    /// is currently unsupported
    fn seqs_to_weights(seqs: &[Vec<u8>], _pseudos: Option<&[f32]>) -> Result<Array2<f32>> {
        let p1 = vec![DEF_PSEUDO; Self::MONO_CT];
        let pseudos = match _pseudos {
            Some(ref p2) => p2,
            None => p1.as_slice(),
        };

        if pseudos.len() != Self::MONO_CT {
            return Err(Error::InvalidPseudos {
                expected: Self::MONO_CT as u8,
                received: pseudos.len() as u8,
            });
        }

        if seqs.is_empty() {
            return Err(Error::EmptyMotif);
        }

        let seqlen = seqs[0].len();
        let mut counts = Array2::zeros((seqlen, Self::MONO_CT));
        for i in 0..seqlen {
            for base in 0..Self::MONO_CT {
                counts[[i, base]] = pseudos[base];
            }
        }

        for seq in seqs.iter() {
            if seq.len() != seqlen {
                return Err(Error::InconsistentLen);
            }

            for (idx, base) in seq.iter().enumerate() {
                match Self::lookup(*base) {
                    Err(e) => return Err(e),
                    Ok(pos) => counts[[idx, pos]] += 1.0,
                }
            }
        }
        Ok(counts)
    }

    /// Returns the index of given monomer in the scores matrix using the lookup table `LK`
    /// # Arguments
    /// * `mono` - monomer, eg, b'A' for DNA or b'R' for protein
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - `mono` wasn't found in the lookup table
    fn lookup(mono: u8) -> Result<usize> {
        if mono >= 127 {
            Err(Error::InvalidMonomer { mono })
        } else {
            let idx = Self::LK[mono as usize];
            if idx == INVALID_MONO {
                Err(Error::InvalidMonomer { mono })
            } else {
                Ok(idx as usize)
            }
        }
    }

    /// Returns the monomer associated with the given index; the reverse of `lookup`.
    /// Returns INVALID_MONO if the index isn't associated with a monomer.
    /// # Arguments
    /// * `idx` - the index in question
    fn rev_lk(idx: usize) -> u8;

    /// Returns the length of motif
    fn len(&self) -> usize;

    /// Returns a representation of the motif using ambiguous codes.
    /// Primarily useful for DNA motifs, where ambiguous codes are
    /// common (eg, 'M' for 'A or C'); less so for proteins, where we
    /// represent any position without a dominant amino acid as an 'X'
    fn degenerate_consensus(&self) -> Vec<u8>;

    /// Accessor - returns scores matrix
    fn get_scores(&self) -> &Array2<f32>;

    /// Return sum of "worst" base at each position
    fn get_min_score(&self) -> f32;

    /// Return sum of "best" base at each position
    fn get_max_score(&self) -> f32;

    /// Returns information content of a single position.
    /// Used `info_content` method.
    /// FIXME: this should be replaced with a CTFE ... or maybe just a constant
    fn get_bits() -> f32;

    /// Returns the un-normalized sum of matching bases, useful for comparing matches from
    /// motifs of different lengths
    ///
    /// # Arguments
    /// * `seq_it` - iterator representing the query sequence
    ///
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - sequence `seq_it` contained invalid monomer `mono`
    fn raw_score<C, T>(&self, seq_it: T) -> Result<(usize, f32, Vec<f32>)>
    where
        C: Borrow<u8>,
        T: IntoIterator<Item = C>,
    {
        let pssm_len = self.len();

        let mut best_start = 0;
        let mut best_score = -1.0;
        let mut best_m = Vec::new();
        // we have to look at slices, so a simple iterator won't do
        let seq = seq_it.into_iter().map(|c| *c.borrow()).collect_vec();
        let scores = self.get_scores();
        for start in 0..=seq.len() - pssm_len {
            let m: Vec<f32> = match (0..pssm_len)
                .map(|i| match Self::lookup(seq[start + i]) {
                    Err(e) => Err(e),
                    Ok(pos) => Ok(scores[[i, pos]]),
                })
                .collect()
            {
                Ok(m) => m,
                Err(e) => return Err(e),
            };
            let tot = m.iter().sum();
            if tot > best_score {
                best_score = tot;
                best_start = start;
                best_m = m;
            }
        }
        Ok((best_start, best_score, best_m))
    }

    /// Returns a `ScoredPos` struct representing the best match within the query sequence
    /// see:
    ///   MATCHTM: a tool for searching transcription factor binding sites in DNA sequences
    ///   Nucleic Acids Res. 2003 Jul 1; 31(13): 3576–3579
    ///   https://www.ncbi.nlm.nih.gov/pmc/articles/PMC169193/
    ///
    /// # Arguments
    /// * `seq_it` - iterator representing the query sequence
    ///
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - sequence `seq_it` contained invalid monomer `mono`
    /// * `Error::QueryTooShort` - sequence `seq_id` was too short
    ///
    /// # Example
    /// let pssm = DNAMotif::from_seqs(vec![
    ///            b"AAAA".to_vec(),
    ///            b"AATA".to_vec(),
    ///            b"AAGA".to_vec(),
    ///            b"AAAA".to_vec(),
    ///        ].as_ref(), None).unwrap();
    /// let start_pos = pssm.score(b"CCCCCAATA").unwrap().loc;
    fn score<C, T>(&self, seq_it: T) -> Result<ScoredPos>
    where
        C: Borrow<u8>,
        T: IntoIterator<Item = C>,
    {
        let pssm_len = self.len();
        let seq = seq_it.into_iter().map(|c| *c.borrow()).collect_vec();
        if seq.len() < pssm_len {
            return Err(Error::QueryTooShort {
                motif_len: pssm_len,
                query_len: seq.len(),
            });
        }
        let min_score = self.get_min_score();
        let max_score = self.get_max_score();

        if approx_eq!(max_score, min_score) {
            return Err(Error::NullMotif);
        }

        let (best_start, best_score, best_m) = self.raw_score(&seq)?;

        Ok(ScoredPos {
            loc: best_start,
            sum: (best_score - min_score) / (max_score - min_score),
            scores: best_m,
        })
    }

    /// Returns a float representing the information content of a motif; roughly the
    /// inverse of Shannon Entropy.
    /// Adapted from the information content described here:
    ///    https://en.wikipedia.org/wiki/Sequence_logo#Logo_creation
    fn info_content(&self) -> f32 {
        fn ent<'a, I>(probs: I) -> f32
        where
            I: Iterator<Item = &'a f32>,
        {
            probs
                .map(|p| {
                    if *p == 0.0 {
                        0.0
                    } else {
                        -1.0 * *p * p.log(2.0)
                    }
                })
                .sum()
        }
        let bits = Self::get_bits();
        let scores = self.get_scores();
        let mut tot = 0.0;
        for row in scores.genrows() {
            tot += bits - ent(row.iter());
        }
        tot
    }
}
<commit_msg>use absolute difference as suggested by clippy<commit_after>// Copyright 2018 Kieran Hervold
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Create a weight matrix representing a set of aligned reference sequences
//! that constitute a motif, and use this matrix to scan query sequences for
//! occurrences of this motif.
//! Complexity: O(n*m) for motif length n and query length m
//!
//! The position-specific scoring matrix (PSSM), aka position weight matrix (PWM),
//! algorithm is implemented for both DNA and amino-acid sequences.
//!
//! # Examples
//!
//! use bio::pattern_matching::pssm::DNAMotif;
//! let pssm = DNAMotif::from_seqs(vec![
//!            b"AAAA".to_vec(),
//!            b"AATA".to_vec(),
//!            b"AAGA".to_vec(),
//!            b"AAAA".to_vec(),
//!        ].as_ref(), None).unwrap();
//! let start_pos = pssm.score(b"CCCCCAATA").unwrap().loc;
//! println!("motif found at position {}", start_pos);
//!
//! /* amino acid sequences are supported, too */
//! use pssm::pattern_matching::pssm::ProtMotif;
//! let pssm = ProtMotif::from_seqs(vec![
//!            b"ARNNYM".to_vec(),
//!            b"ARNRYM".to_vec(),
//!            b"ARNNCM".to_vec(),
//!            b"ARNNYM".to_vec(),
//!        ].as_ref(), None).unwrap();

use std::borrow::Borrow;
use std::f32::NEG_INFINITY;

use itertools::Itertools;
use ndarray::prelude::Array2;

mod dnamotif;
pub mod errors;
mod protmotif;

pub use self::dnamotif::DNAMotif;
pub use self::errors::{Error, Result};
pub use self::protmotif::ProtMotif;

/// default pseudocount - used to prevent 0 tallies
pub const DEF_PSEUDO: f32 = 0.5;
/// approximately zero
pub const EPSILON: f32 = 1e-5;
/// value representing an invalid monomer in lookup table
pub const INVALID_MONO: u8 = 255;

/// Represents motif score & location of match
#[derive(Debug, Clone, PartialEq)]
pub struct ScoredPos {
    pub loc: usize,
    pub sum: f32,
    pub scores: Vec<f32>,
}

impl Default for ScoredPos {
    fn default() -> ScoredPos {
        ScoredPos {
            loc: 0,
            sum: NEG_INFINITY,
            scores: Vec::new(),
        }
    }
}

/// Trait containing code shared between DNA and protein implementations
/// of the position-specific scoring matrix.
pub trait Motif {
    /// Lookup table mapping monomer -> index
    const LK: [u8; 127] = [INVALID_MONO; 127];
    /// All monomers, in order corresponding to lookup table
    const MONOS: &'static [u8] = b"";
    /// Monomer count - equal to length of `MONOS`
    const MONO_CT: usize = 0;

    /// Returns a weight matrix representing the sequences provided.
    /// This code is shared by implementations of `from_seqs`
    /// # Arguments
    /// * `seqs` - sequences incorporated into motif
    /// * `pseudos` - array slice with a pseudocount for each monomer;
    ///    defaults to DEF_PSEUDO for all if None is supplied
    ///
    /// FIXME: pseudos should be an array of size MONO_CT, but that
    /// is currently unsupported
    fn seqs_to_weights(seqs: &[Vec<u8>], _pseudos: Option<&[f32]>) -> Result<Array2<f32>> {
        let p1 = vec![DEF_PSEUDO; Self::MONO_CT];
        let pseudos = match _pseudos {
            Some(ref p2) => p2,
            None => p1.as_slice(),
        };

        if pseudos.len() != Self::MONO_CT {
            return Err(Error::InvalidPseudos {
                expected: Self::MONO_CT as u8,
                received: pseudos.len() as u8,
            });
        }

        if seqs.is_empty() {
            return Err(Error::EmptyMotif);
        }

        let seqlen = seqs[0].len();
        let mut counts = Array2::zeros((seqlen, Self::MONO_CT));
        for i in 0..seqlen {
            for base in 0..Self::MONO_CT {
                counts[[i, base]] = pseudos[base];
            }
        }

        for seq in seqs.iter() {
            if seq.len() != seqlen {
                return Err(Error::InconsistentLen);
            }

            for (idx, base) in seq.iter().enumerate() {
                match Self::lookup(*base) {
                    Err(e) => return Err(e),
                    Ok(pos) => counts[[idx, pos]] += 1.0,
                }
            }
        }
        Ok(counts)
    }

    /// Returns the index of given monomer in the scores matrix using the lookup table `LK`
    /// # Arguments
    /// * `mono` - monomer, eg, b'A' for DNA or b'R' for protein
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - `mono` wasn't found in the lookup table
    fn lookup(mono: u8) -> Result<usize> {
        if mono >= 127 {
            Err(Error::InvalidMonomer { mono })
        } else {
            let idx = Self::LK[mono as usize];
            if idx == INVALID_MONO {
                Err(Error::InvalidMonomer { mono })
            } else {
                Ok(idx as usize)
            }
        }
    }

    /// Returns the monomer associated with the given index; the reverse of `lookup`.
    /// Returns INVALID_MONO if the index isn't associated with a monomer.
    /// # Arguments
    /// * `idx` - the index in question
    fn rev_lk(idx: usize) -> u8;

    /// Returns the length of motif
    fn len(&self) -> usize;

    /// Returns a representation of the motif using ambiguous codes.
    /// Primarily useful for DNA motifs, where ambiguous codes are
    /// common (eg, 'M' for 'A or C'); less so for proteins, where we
    /// represent any position without a dominant amino acid as an 'X'
    fn degenerate_consensus(&self) -> Vec<u8>;

    /// Accessor - returns scores matrix
    fn get_scores(&self) -> &Array2<f32>;

    /// Return sum of "worst" base at each position
    fn get_min_score(&self) -> f32;

    /// Return sum of "best" base at each position
    fn get_max_score(&self) -> f32;

    /// Returns information content of a single position.
    /// Used `info_content` method.
    /// FIXME: this should be replaced with a CTFE ... or maybe just a constant
    fn get_bits() -> f32;

    /// Returns the un-normalized sum of matching bases, useful for comparing matches from
    /// motifs of different lengths
    ///
    /// # Arguments
    /// * `seq_it` - iterator representing the query sequence
    ///
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - sequence `seq_it` contained invalid monomer `mono`
    fn raw_score<C, T>(&self, seq_it: T) -> Result<(usize, f32, Vec<f32>)>
    where
        C: Borrow<u8>,
        T: IntoIterator<Item = C>,
    {
        let pssm_len = self.len();

        let mut best_start = 0;
        let mut best_score = -1.0;
        let mut best_m = Vec::new();
        // we have to look at slices, so a simple iterator won't do
        let seq = seq_it.into_iter().map(|c| *c.borrow()).collect_vec();
        let scores = self.get_scores();
        for start in 0..=seq.len() - pssm_len {
            let m: Vec<f32> = match (0..pssm_len)
                .map(|i| match Self::lookup(seq[start + i]) {
                    Err(e) => Err(e),
                    Ok(pos) => Ok(scores[[i, pos]]),
                })
                .collect()
            {
                Ok(m) => m,
                Err(e) => return Err(e),
            };
            let tot = m.iter().sum();
            if tot > best_score {
                best_score = tot;
                best_start = start;
                best_m = m;
            }
        }
        Ok((best_start, best_score, best_m))
    }

    /// Returns a `ScoredPos` struct representing the best match within the query sequence
    /// see:
    ///   MATCHTM: a tool for searching transcription factor binding sites in DNA sequences
    ///   Nucleic Acids Res. 2003 Jul 1; 31(13): 3576–3579
    ///   https://www.ncbi.nlm.nih.gov/pmc/articles/PMC169193/
    ///
    /// # Arguments
    /// * `seq_it` - iterator representing the query sequence
    ///
    /// # Errors
    /// * `Error::InvalidMonomer(mono)` - sequence `seq_it` contained invalid monomer `mono`
    /// * `Error::QueryTooShort` - sequence `seq_id` was too short
    ///
    /// # Example
    /// let pssm = DNAMotif::from_seqs(vec![
    ///            b"AAAA".to_vec(),
    ///            b"AATA".to_vec(),
    ///            b"AAGA".to_vec(),
    ///            b"AAAA".to_vec(),
    ///        ].as_ref(), None).unwrap();
    /// let start_pos = pssm.score(b"CCCCCAATA").unwrap().loc;
    fn score<C, T>(&self, seq_it: T) -> Result<ScoredPos>
    where
        C: Borrow<u8>,
        T: IntoIterator<Item = C>,
    {
        let pssm_len = self.len();
        let seq = seq_it.into_iter().map(|c| *c.borrow()).collect_vec();
        if seq.len() < pssm_len {
            return Err(Error::QueryTooShort {
                motif_len: pssm_len,
                query_len: seq.len(),
            });
        }
        let min_score = self.get_min_score();
        let max_score = self.get_max_score();

        if abs_diff_eq!(max_score, min_score) {
            return Err(Error::NullMotif);
        }

        let (best_start, best_score, best_m) = self.raw_score(&seq)?;

        Ok(ScoredPos {
            loc: best_start,
            sum: (best_score - min_score) / (max_score - min_score),
            scores: best_m,
        })
    }

    /// Returns a float representing the information content of a motif; roughly the
    /// inverse of Shannon Entropy.
    /// Adapted from the information content described here:
    ///    https://en.wikipedia.org/wiki/Sequence_logo#Logo_creation
    fn info_content(&self) -> f32 {
        fn ent<'a, I>(probs: I) -> f32
        where
            I: Iterator<Item = &'a f32>,
        {
            probs
                .map(|p| {
                    if *p == 0.0 {
                        0.0
                    } else {
                        -1.0 * *p * p.log(2.0)
                    }
                })
                .sum()
        }
        let bits = Self::get_bits();
        let scores = self.get_scores();
        let mut tot = 0.0;
        for row in scores.genrows() {
            tot += bits - ent(row.iter());
        }
        tot
    }
}
<|endoftext|>