<commit_before>//! Singular-value decomposition (SVD) by divide-and-conquer (?gesdd)

use ndarray::*;

use super::error::*;
use super::layout::*;
use super::types::*;

pub use lapack::svddc::UVTFlag;

/// Singular-value decomposition of matrix (copying) by divide-and-conquer
pub trait SVDDC {
    type U;
    type VT;
    type Sigma;
    fn svddc(&self, uvt_flag: UVTFlag) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

/// Singular-value decomposition of matrix by divide-and-conquer
pub trait SVDDCInto {
    type U;
    type VT;
    type Sigma;
    fn svddc_into(
        self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

/// Singular-value decomposition of matrix reference by divide-and-conquer
pub trait SVDDCInplace {
    type U;
    type VT;
    type Sigma;
    fn svddc_inplace(
        &mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

impl<A, S> SVDDC for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc(&self, uvt_flag: UVTFlag) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        self.to_owned().svddc_into(uvt_flag)
    }
}

impl<A, S> SVDDCInto for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc_into(
        mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        self.svddc_inplace(uvt_flag)
    }
}

impl<A, S> SVDDCInplace for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc_inplace(
        &mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        let l = self.layout()?;
        let svd_res = unsafe { A::svddc(l, uvt_flag, self.as_allocated_mut()?)? };
        let (n, m) = l.size();
        let k = std::cmp::min(n, m);
        let n = n as usize;
        let m = m as usize;
        let k = k as usize;

        let (u_col, vt_row) = match uvt_flag {
            UVTFlag::Full => (n, m),
            UVTFlag::Some => (k, k),
            UVTFlag::None => (0, 0),
        };

        let u = svd_res.u.map(|u| {
            assert_eq!(u.len(), n * u_col);
            match l {
                MatrixLayout::F { .. } => Array::from_shape_vec((n, u_col).f(), u),
                MatrixLayout::C { .. } => Array::from_shape_vec((n, u_col), u),
            }
            .unwrap()
        });

        let vt = svd_res.vt.map(|vt| {
            assert_eq!(vt.len(), m * vt_row);
            match l {
                MatrixLayout::F { .. } => Array::from_shape_vec((vt_row, m).f(), vt),
                MatrixLayout::C { .. } => Array::from_shape_vec((vt_row, m), vt),
            }
            .unwrap()
        });

        let s = ArrayBase::from(svd_res.s);
        Ok((u, s, vt))
    }
}
<commit_msg>Use convert::into_matrix<commit_after>//! Singular-value decomposition (SVD) by divide-and-conquer (?gesdd)

use super::{convert::*, error::*, layout::*, types::*};
use ndarray::*;

pub use lapack::svddc::UVTFlag;

/// Singular-value decomposition of matrix (copying) by divide-and-conquer
pub trait SVDDC {
    type U;
    type VT;
    type Sigma;
    fn svddc(&self, uvt_flag: UVTFlag) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

/// Singular-value decomposition of matrix by divide-and-conquer
pub trait SVDDCInto {
    type U;
    type VT;
    type Sigma;
    fn svddc_into(
        self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

/// Singular-value decomposition of matrix reference by divide-and-conquer
pub trait SVDDCInplace {
    type U;
    type VT;
    type Sigma;
    fn svddc_inplace(
        &mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)>;
}

impl<A, S> SVDDC for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc(&self, uvt_flag: UVTFlag) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        self.to_owned().svddc_into(uvt_flag)
    }
}

impl<A, S> SVDDCInto for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc_into(
        mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        self.svddc_inplace(uvt_flag)
    }
}

impl<A, S> SVDDCInplace for ArrayBase<S, Ix2>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
{
    type U = Array2<A>;
    type VT = Array2<A>;
    type Sigma = Array1<A::Real>;

    fn svddc_inplace(
        &mut self,
        uvt_flag: UVTFlag,
    ) -> Result<(Option<Self::U>, Self::Sigma, Option<Self::VT>)> {
        let l = self.layout()?;
        let svd_res = unsafe { A::svddc(l, uvt_flag, self.as_allocated_mut()?)? };
        let (m, n) = l.size();
        let k = m.min(n);

        let (u_col, vt_row) = match uvt_flag {
            UVTFlag::Full => (m, n),
            UVTFlag::Some => (k, k),
            UVTFlag::None => (0, 0),
        };

        let u = svd_res
            .u
            .map(|u| into_matrix(l.resized(m, u_col), u).unwrap());

        let vt = svd_res
            .vt
            .map(|vt| into_matrix(l.resized(vt_row, n), vt).unwrap());

        let s = ArrayBase::from(svd_res.s);
        Ok((u, s, vt))
    }
}
<|endoftext|>


<commit_before>use super::*;

///
/// A position in a source string
///
pub struct Position {
    line: uint,
    col: uint,
}
static DUMMY_POS: Position = Position { line: 0, col: 0 };

pub struct ParseFailure {
    description: ~str,
    position: Position,
}

struct Token<'self> {
    val: &'self str,
    position: Position,
}

impl<'self> Token<'self> {
    fn new<'r>(val: &'r str, position: Position) -> Token<'r> {
        Token { val: val, position: position }
    }
}

pub struct Parser<'self> {
    token_start: uint,
    position: uint,
    token: Option<Token<'self>>,
    src: &'self str
}

impl<'self> Parser<'self> {
    fn new<'r>(src: &'r str) -> Parser<'r> {
        Parser {
            token_start: 0,
            position: 0,
            token: None,
            src: src
        }
    }

    // Tokenising

    fn load_token(&mut self) {
        cond!(
            // already loaded
            (self.token.is_some()) { return; }
            // a left-over token
            (self.token_start < self.position) {
                self.token = Some(Token::new(self.src.slice(self.token_start, self.position),
                                             DUMMY_POS));
                self.token_start = self.position;
                return;
            }
            // finished
            (self.token_start >= self.src.len()) { self.token = None; return; }
        );

        let mut seen_token = false;
        for str::each_char(self.src.slice(self.position, self.src.len())) |c| {
            // TODO: tokenise strings
            match c {
                '(' | ')' | '|' => {
                    if !seen_token {
                        // we haven't seen a token yet, so use this character
                        self.token_start = self.position;
                        self.position += 1; // always 1 byte
                        seen_token = true;
                    } // we've already seen a token, so don't advance, and use that one
                    break;
                }
                // ignore (leading) whitespace
                _ if c.is_whitespace() => { if seen_token { break } }
                _ => {
                    if !seen_token { // start a new token here
                        self.token_start = self.position;
                        seen_token = true;
                    }
                }
            }
            self.position += char::len_utf8_bytes(c);
        }
        self.token = if seen_token {
            Some(Token::new(self.src.slice(self.token_start, self.position), DUMMY_POS))
        } else {
            None
        };
        self.token_start = self.position;
    }

    fn peek_token(&mut self) -> Option<Token<'self>> {
        // this is does nothing if token is Some
        self.load_token();
        self.token
    }

    fn peek_token_no_eof(&mut self) -> Result<Token<'self>, ParseFailure> {
        match self.peek_token() {
            Some(tok) => Ok(tok),
            None => Err(ParseFailure { position: DUMMY_POS, description: ~"Unexpected EOF"})
        }
    }

    fn bump_token(&mut self) {
        self.load_token();
        self.token = None;
    }

    fn eat_token(&mut self, tok: &str) -> bool {
        match self.peek_token() {
            Some(Token { val, _ }) if val == tok => {
                self.bump_token();
                true
            }
            _ => false,
        }
    }

    fn expect_token(&mut self, tok: &str) -> Result<(), ParseFailure> {
        if self.eat_token(tok) {
            Ok(())
        } else {
            Err(ParseFailure { position: DUMMY_POS, description: fmt!("Expecting %s", tok)})
        }
    }

    // Parsing

    /// Parse an identifier
    fn parse_ident(&mut self) -> Result<Ident, ParseFailure> {
        do self.peek_token_no_eof().chain |tok| {
            match tok.val {
                "" | "|" | "(" | ")" => {
                    Err(ParseFailure { position: tok.position, description: ~"empty ident" })
                }
                _ => {
                    self.bump_token();
                    Ok(Ident(tok.val.to_owned()))
                }
            }
        }
    }

    fn parse_value_from_token(&mut self, tok: Token)  -> Result<Value, ParseFailure> {
        macro_rules! ret(($val:expr) => {{ self.bump_token(); return Ok($val); }});
        match tok.val {
            "true" => ret!(Boolean(true)),
            "false" => ret!(Boolean(false)),
            _ => {}
        }
        match int::from_str(tok.val) {
            Some(i) => ret!(Integer(i)),
            None => {}
        }
        match float::from_str(tok.val) {
            Some(f) => ret!(Float(f)),
            None => {}
        }

        Err(ParseFailure { position: tok.position, description: ~"Invalid value"})
    }

    fn parse_value(&mut self) -> Result<Value, ParseFailure> {
        do self.peek_token_no_eof().chain |tok| { self.parse_value_from_token(tok) }
    }

    fn parse_value_or_ident(&mut self) -> Result<@Expr, ParseFailure> {
        match self.parse_value() {
            Ok(val) => Ok(@Literal(val)),
            Err(_) => do self.parse_ident().map |&ident| { @Symbol(ident) }
        }
    }

    /// Parse an if expression, without the leading 'if'
    fn parse_if(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().chain |pred| {
            do self.parse().chain |then| {
                do self.parse().map |&els| {
                    @IfExpr(pred, then, els)
                }
            }
        }
    }

    /// Parse a lambda expression, without the leading 'fn'
    fn parse_lambda(&mut self) -> Result<@Expr, ParseFailure> {
        do self.expect_token("(").chain |_| {
            let args = do vec::build |push| {
                loop {
                    match self.parse_ident() {
                        Ok(ident) => push(ident),
                        Err(_) => break,
                    }
                }
            };

            match self.expect_token(")") {
                Err(err) => Err(err),
                _ => match self.parse() {
                    Ok(expr) => Ok(@Literal(Lambda(args, expr, Env::empty()))),
                    Err(err) => Err(err)
                }
            }
        }
    }

    /// Parse a quote expression, without the leading 'quote'
    fn parse_quote(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().map |&expr| {
            @Literal(Quote(expr))
        }
    }

    fn parse_def(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse_ident().chain |ident| {
            do self.parse().map |&expr| {
                @LetExpr(copy ident, expr)
            }
        }
    }

    fn parse_call(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().chain |expr| {
            let args = do vec::build |push| {
                loop {
                    match self.parse() {
                        Ok(expr) => push(expr),
                        Err(_) => break
                    }
                }
            };

            Ok(@CallExpr(expr, args))
        }
    }

    /// Parse the interior of an S-expr
    fn parse_parened(&mut self) -> Result<@Expr, ParseFailure> {
        cond! (
            (self.eat_token("if"))    { self.parse_if() }
            (self.eat_token("quote")) { self.parse_quote() }
            (self.eat_token("let"))   { self.parse_def() }
            (self.eat_token("fn"))    { self.parse_lambda() }
            _ { self.parse_call() }
        )
    }

    fn parse(&mut self) -> Result<@Expr, ParseFailure> {
        if self.eat_token("(") {
            let ret = self.parse_parened();
            if !self.eat_token(")") {
                Err(ParseFailure { position: DUMMY_POS, description: ~"Expecting ')'"})
            } else {
                ret
            }
        } else {
            do self.peek_token_no_eof().chain |tok| {
                if (tok.val == ")") {
                    Err(ParseFailure { position: tok.position, description: ~"Unexpected ')'"})
                } else {
                    self.parse_value_or_ident()
                }
            }
        }
    }
}

///
/// Performs a recursive decent parse of the source string.
///
pub fn parse(src: &str) -> Result<@Expr, ParseFailure> {
    Parser::new(src).parse()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parser_ok() {
        fn test(s: ~str) {
            assert_eq!(Parser::new(s).parse().get().to_str(),
                       s);
        }

        test(~"1");
        test(~"(let a (+ 1 2))");

        // the extra space after the lambda?
        test(~"(if true (fn (a b) (+ 1 a b)) (quote (1 2 3)))");
    }

    #[test]
    fn test_parser_err() {
        fn test(s: ~str) {
             assert!(Parser::new(s).parse().is_err());
        }

        test(~"(");
        test(~")");

        test(~"(fn (a b (+ 1 a b)))");
        test(~"(if true)");
        test(~"(if true 1 2 3)");
        test(~"(let a)");
        test(~"(let a b c)");
    }
}<commit_msg>Remove stray pipe tests<commit_after>use super::*;

///
/// A position in a source string
///
pub struct Position {
    line: uint,
    col: uint,
}
static DUMMY_POS: Position = Position { line: 0, col: 0 };

pub struct ParseFailure {
    description: ~str,
    position: Position,
}

struct Token<'self> {
    val: &'self str,
    position: Position,
}

impl<'self> Token<'self> {
    fn new<'r>(val: &'r str, position: Position) -> Token<'r> {
        Token { val: val, position: position }
    }
}

pub struct Parser<'self> {
    token_start: uint,
    position: uint,
    token: Option<Token<'self>>,
    src: &'self str
}

impl<'self> Parser<'self> {
    fn new<'r>(src: &'r str) -> Parser<'r> {
        Parser {
            token_start: 0,
            position: 0,
            token: None,
            src: src
        }
    }

    // Tokenising

    fn load_token(&mut self) {
        cond!(
            // already loaded
            (self.token.is_some()) { return; }
            // a left-over token
            (self.token_start < self.position) {
                self.token = Some(Token::new(self.src.slice(self.token_start, self.position),
                                             DUMMY_POS));
                self.token_start = self.position;
                return;
            }
            // finished
            (self.token_start >= self.src.len()) { self.token = None; return; }
        );

        let mut seen_token = false;
        for str::each_char(self.src.slice(self.position, self.src.len())) |c| {
            // TODO: tokenise strings
            match c {
                '(' | ')' => {
                    if !seen_token {
                        // we haven't seen a token yet, so use this character
                        self.token_start = self.position;
                        self.position += 1; // always 1 byte
                        seen_token = true;
                    } // we've already seen a token, so don't advance, and use that one
                    break;
                }
                // ignore (leading) whitespace
                _ if c.is_whitespace() => { if seen_token { break } }
                _ => {
                    if !seen_token { // start a new token here
                        self.token_start = self.position;
                        seen_token = true;
                    }
                }
            }
            self.position += char::len_utf8_bytes(c);
        }
        self.token = if seen_token {
            Some(Token::new(self.src.slice(self.token_start, self.position), DUMMY_POS))
        } else {
            None
        };
        self.token_start = self.position;
    }

    fn peek_token(&mut self) -> Option<Token<'self>> {
        // this is does nothing if token is Some
        self.load_token();
        self.token
    }

    fn peek_token_no_eof(&mut self) -> Result<Token<'self>, ParseFailure> {
        match self.peek_token() {
            Some(tok) => Ok(tok),
            None => Err(ParseFailure { position: DUMMY_POS, description: ~"Unexpected EOF"})
        }
    }

    fn bump_token(&mut self) {
        self.load_token();
        self.token = None;
    }

    fn eat_token(&mut self, tok: &str) -> bool {
        match self.peek_token() {
            Some(Token { val, _ }) if val == tok => {
                self.bump_token();
                true
            }
            _ => false,
        }
    }

    fn expect_token(&mut self, tok: &str) -> Result<(), ParseFailure> {
        if self.eat_token(tok) {
            Ok(())
        } else {
            Err(ParseFailure { position: DUMMY_POS, description: fmt!("Expecting %s", tok)})
        }
    }

    // Parsing

    /// Parse an identifier
    fn parse_ident(&mut self) -> Result<Ident, ParseFailure> {
        do self.peek_token_no_eof().chain |tok| {
            match tok.val {
                "" | "(" | ")" => {
                    Err(ParseFailure { position: tok.position, description: ~"empty ident" })
                }
                _ => {
                    self.bump_token();
                    Ok(Ident(tok.val.to_owned()))
                }
            }
        }
    }

    fn parse_value_from_token(&mut self, tok: Token)  -> Result<Value, ParseFailure> {
        macro_rules! ret(($val:expr) => {{ self.bump_token(); return Ok($val); }});
        match tok.val {
            "true" => ret!(Boolean(true)),
            "false" => ret!(Boolean(false)),
            _ => {}
        }
        match int::from_str(tok.val) {
            Some(i) => ret!(Integer(i)),
            None => {}
        }
        match float::from_str(tok.val) {
            Some(f) => ret!(Float(f)),
            None => {}
        }

        Err(ParseFailure { position: tok.position, description: ~"Invalid value"})
    }

    fn parse_value(&mut self) -> Result<Value, ParseFailure> {
        do self.peek_token_no_eof().chain |tok| { self.parse_value_from_token(tok) }
    }

    fn parse_value_or_ident(&mut self) -> Result<@Expr, ParseFailure> {
        match self.parse_value() {
            Ok(val) => Ok(@Literal(val)),
            Err(_) => do self.parse_ident().map |&ident| { @Symbol(ident) }
        }
    }

    /// Parse an if expression, without the leading 'if'
    fn parse_if(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().chain |pred| {
            do self.parse().chain |then| {
                do self.parse().map |&els| {
                    @IfExpr(pred, then, els)
                }
            }
        }
    }

    /// Parse a lambda expression, without the leading 'fn'
    fn parse_lambda(&mut self) -> Result<@Expr, ParseFailure> {
        do self.expect_token("(").chain |_| {
            let args = do vec::build |push| {
                loop {
                    match self.parse_ident() {
                        Ok(ident) => push(ident),
                        Err(_) => break,
                    }
                }
            };

            match self.expect_token(")") {
                Err(err) => Err(err),
                _ => match self.parse() {
                    Ok(expr) => Ok(@Literal(Lambda(args, expr, Env::empty()))),
                    Err(err) => Err(err)
                }
            }
        }
    }

    /// Parse a quote expression, without the leading 'quote'
    fn parse_quote(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().map |&expr| {
            @Literal(Quote(expr))
        }
    }

    fn parse_def(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse_ident().chain |ident| {
            do self.parse().map |&expr| {
                @LetExpr(copy ident, expr)
            }
        }
    }

    fn parse_call(&mut self) -> Result<@Expr, ParseFailure> {
        do self.parse().chain |expr| {
            let args = do vec::build |push| {
                loop {
                    match self.parse() {
                        Ok(expr) => push(expr),
                        Err(_) => break
                    }
                }
            };

            Ok(@CallExpr(expr, args))
        }
    }

    /// Parse the interior of an S-expr
    fn parse_parened(&mut self) -> Result<@Expr, ParseFailure> {
        cond! (
            (self.eat_token("if"))    { self.parse_if() }
            (self.eat_token("quote")) { self.parse_quote() }
            (self.eat_token("let"))   { self.parse_def() }
            (self.eat_token("fn"))    { self.parse_lambda() }
            _ { self.parse_call() }
        )
    }

    fn parse(&mut self) -> Result<@Expr, ParseFailure> {
        if self.eat_token("(") {
            let ret = self.parse_parened();
            if !self.eat_token(")") {
                Err(ParseFailure { position: DUMMY_POS, description: ~"Expecting ')'"})
            } else {
                ret
            }
        } else {
            do self.peek_token_no_eof().chain |tok| {
                if (tok.val == ")") {
                    Err(ParseFailure { position: tok.position, description: ~"Unexpected ')'"})
                } else {
                    self.parse_value_or_ident()
                }
            }
        }
    }
}

///
/// Performs a recursive decent parse of the source string.
///
pub fn parse(src: &str) -> Result<@Expr, ParseFailure> {
    Parser::new(src).parse()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parser_ok() {
        fn test(s: ~str) {
            assert_eq!(Parser::new(s).parse().get().to_str(),
                       s);
        }

        test(~"1");
        test(~"(let a (+ 1 2))");

        // the extra space after the lambda?
        test(~"(if true (fn (a b) (+ 1 a b)) (quote (1 2 3)))");
    }

    #[test]
    fn test_parser_err() {
        fn test(s: ~str) {
             assert!(Parser::new(s).parse().is_err());
        }

        test(~"(");
        test(~")");

        test(~"(fn (a b (+ 1 a b)))");
        test(~"(if true)");
        test(~"(if true 1 2 3)");
        test(~"(let a)");
        test(~"(let a b c)");
    }
}<|endoftext|>


<commit_before>use index::btree::internal::InNode;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::BPlusTree;
use index::btree::NodeCellRef;
use index::lsmtree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::mem;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

enum PruningSearch {
    DeepestInnode,
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn apply_removal<'a, KS, PS>(
    cursor_guard: &mut NodeWriteGuard<KS, PS>,
    poses: &mut BTreeSet<usize>,
    empty_pages: &mut Vec<&'a EntryKey>,
    prev_key: &Option<&'a EntryKey>,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    if poses.is_empty() {
        return;
    }
    debug!(
        "Applying removal, guard key {:?}, poses {:?}",
        cursor_guard.first_key(),
        poses
    );
    {
        let mut innode = cursor_guard.innode_mut();
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            let mut keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, k)| k)
                .collect();
            let mut ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, p)| p)
                .collect();
            innode.len = keys.len();
            debug!("Prune filtered page have keys {:?}", &keys);
            for (i, key) in keys.into_iter().enumerate() {
                mem::swap(&mut new_keys.as_slice()[i], key);
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                mem::swap(&mut new_ptrs.as_slice()[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;

        debug!(
            "Pruned page have keys {:?}",
            &innode.keys.as_slice_immute()[..innode.len]
        );
    }

    if cursor_guard.is_empty() {
        debug!("Pruned page is empty: {:?}", prev_key);
        cursor_guard.make_empty_node();
        if let &Some(k) = prev_key {
            empty_pages.push(k);
        }
    }
    poses.clear();
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, mut keys: Vec<&'a EntryKey>, level: usize) -> Vec<&'a EntryKey>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let key_len = keys.len();
    let first_search = mut_search::<KS, PS>(node, keys.first().unwrap());
    let pruning = match first_search {
        MutSearchResult::Internal(sub_node) => {
            if read_unchecked::<KS, PS>(&sub_node).is_ext() {
                PruningSearch::DeepestInnode
            } else {
                PruningSearch::Innode(sub_node)
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    match pruning {
        PruningSearch::DeepestInnode => {
            debug!("Removing in deepest nodes keys {:?}, level {}", &keys, level);
        }
        PruningSearch::Innode(sub_node) => {
            keys = prune_selected::<KS, PS>(&sub_node, keys, level + 1);
        }
    }
    // empty page references that will dealt with by upper level
    let mut empty_pages = vec![];
    if !keys.is_empty() {
        debug!("Pruning page containing keys {:?}, level {}", &keys, level);
        // start delete
        let mut cursor_guard = write_node::<KS, PS>(node);
        let mut guard_removing_poses = BTreeSet::new();
        let mut prev_key = None;
        debug!(
            "Prune deepest node starting key: {:?}, level {}",
            cursor_guard.first_key(), level
        );
        for (i, key_to_del) in keys.into_iter().enumerate() {
            if key_to_del >= cursor_guard.right_bound() {
                debug!(
                    "Applying removal for overflow current page ({}/{}) key: {:?} >= bound: {:?}. guard keys: {:?}, level {}",
                    guard_removing_poses.len(),
                    cursor_guard.len() + 1,
                    key_to_del,
                    cursor_guard.right_bound(),
                    &cursor_guard.keys()[..cursor_guard.len()],
                    level
                );
                apply_removal(
                    &mut cursor_guard,
                    &mut guard_removing_poses,
                    &mut empty_pages,
                    &prev_key,
                );
                cursor_guard = write_targeted(cursor_guard, key_to_del);
                debug_assert!(!cursor_guard.is_empty_node());
                debug!(
                    "Applied removal for overflow current page ({}), level {}",
                    cursor_guard.len(), level
                );
            }
            let pos = cursor_guard.search(key_to_del);
            debug!("Key to delete have position {}, key: {:?}, level {}", pos, key_to_del, level);
            guard_removing_poses.insert(pos);
            prev_key = Some(key_to_del)
        }
        if !guard_removing_poses.is_empty() {
            debug!("Applying removal for last keys {:?}, level {}", &guard_removing_poses, level);
            apply_removal(
                &mut cursor_guard,
                &mut guard_removing_poses,
                &mut empty_pages,
                &prev_key,
            );
        }
    }
    debug!("Have empty nodes {:?}, level {:?}", &empty_pages, level);
    empty_pages
}

pub fn level_merge<KSA, PSA, KSB, PSB>(
    src_tree: &BPlusTree<KSA, PSA>,
    dest_tree: &BPlusTree<KSB, PSB>,
) -> usize
where
    KSA: Slice<EntryKey> + Debug + 'static,
    PSA: Slice<NodeCellRef> + 'static,
    KSB: Slice<EntryKey> + Debug + 'static,
    PSB: Slice<NodeCellRef> + 'static,
{
    let left_most_leaf_guards = select::<KSA, PSA>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    debug!("Merge selected {} pages", left_most_leaf_guards.len());

    // merge to dest_tree
    {
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .cloned()
            .collect_vec();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(keys);
    }

    // cleanup upper level references
    {
        let page_keys = left_most_leaf_guards
            .iter()
            .filter(|g| !g.is_empty())
            .map(|g| g.first_key())
            .collect_vec();
        prune_selected::<KSA, PSA>(&src_tree.get_root(), page_keys, 0);
    }

    // adjust leaf left, right references
    let right_right_most = left_most_leaf_guards
        .last()
        .unwrap()
        .right_ref()
        .unwrap()
        .clone();
    let left_left_most = left_most_leaf_guards
        .first()
        .unwrap()
        .left_ref()
        .unwrap()
        .clone();
    for mut g in left_most_leaf_guards {
        *g = NodeData::Empty(box EmptyNode {
            left: Some(left_left_most.clone()),
            right: right_right_most.clone(),
        })
    }
    merge_page_len
}
<commit_msg>cargo fmt<commit_after>use index::btree::internal::InNode;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::BPlusTree;
use index::btree::NodeCellRef;
use index::lsmtree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::mem;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

enum PruningSearch {
    DeepestInnode,
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn apply_removal<'a, KS, PS>(
    cursor_guard: &mut NodeWriteGuard<KS, PS>,
    poses: &mut BTreeSet<usize>,
    empty_pages: &mut Vec<&'a EntryKey>,
    prev_key: &Option<&'a EntryKey>,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    if poses.is_empty() {
        return;
    }
    debug!(
        "Applying removal, guard key {:?}, poses {:?}",
        cursor_guard.first_key(),
        poses
    );
    {
        let mut innode = cursor_guard.innode_mut();
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            let mut keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, k)| k)
                .collect();
            let mut ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, p)| p)
                .collect();
            innode.len = keys.len();
            debug!("Prune filtered page have keys {:?}", &keys);
            for (i, key) in keys.into_iter().enumerate() {
                mem::swap(&mut new_keys.as_slice()[i], key);
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                mem::swap(&mut new_ptrs.as_slice()[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;

        debug!(
            "Pruned page have keys {:?}",
            &innode.keys.as_slice_immute()[..innode.len]
        );
    }

    if cursor_guard.is_empty() {
        debug!("Pruned page is empty: {:?}", prev_key);
        cursor_guard.make_empty_node();
        if let &Some(k) = prev_key {
            empty_pages.push(k);
        }
    }
    poses.clear();
}

fn prune_selected<'a, KS, PS>(
    node: &NodeCellRef,
    mut keys: Vec<&'a EntryKey>,
    level: usize,
) -> Vec<&'a EntryKey>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let key_len = keys.len();
    let first_search = mut_search::<KS, PS>(node, keys.first().unwrap());
    let pruning = match first_search {
        MutSearchResult::Internal(sub_node) => {
            if read_unchecked::<KS, PS>(&sub_node).is_ext() {
                PruningSearch::DeepestInnode
            } else {
                PruningSearch::Innode(sub_node)
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    match pruning {
        PruningSearch::DeepestInnode => {
            debug!(
                "Removing in deepest nodes keys {:?}, level {}",
                &keys, level
            );
        }
        PruningSearch::Innode(sub_node) => {
            keys = prune_selected::<KS, PS>(&sub_node, keys, level + 1);
        }
    }
    // empty page references that will dealt with by upper level
    let mut empty_pages = vec![];
    if !keys.is_empty() {
        debug!("Pruning page containing keys {:?}, level {}", &keys, level);
        // start delete
        let mut cursor_guard = write_node::<KS, PS>(node);
        let mut guard_removing_poses = BTreeSet::new();
        let mut prev_key = None;
        debug!(
            "Prune deepest node starting key: {:?}, level {}",
            cursor_guard.first_key(),
            level
        );
        for (i, key_to_del) in keys.into_iter().enumerate() {
            if key_to_del >= cursor_guard.right_bound() {
                debug!(
                    "Applying removal for overflow current page ({}/{}) key: {:?} >= bound: {:?}. guard keys: {:?}, level {}",
                    guard_removing_poses.len(),
                    cursor_guard.len() + 1,
                    key_to_del,
                    cursor_guard.right_bound(),
                    &cursor_guard.keys()[..cursor_guard.len()],
                    level
                );
                apply_removal(
                    &mut cursor_guard,
                    &mut guard_removing_poses,
                    &mut empty_pages,
                    &prev_key,
                );
                cursor_guard = write_targeted(cursor_guard, key_to_del);
                debug_assert!(!cursor_guard.is_empty_node());
                debug!(
                    "Applied removal for overflow current page ({}), level {}",
                    cursor_guard.len(),
                    level
                );
            }
            let pos = cursor_guard.search(key_to_del);
            debug!(
                "Key to delete have position {}, key: {:?}, level {}",
                pos, key_to_del, level
            );
            guard_removing_poses.insert(pos);
            prev_key = Some(key_to_del)
        }
        if !guard_removing_poses.is_empty() {
            debug!(
                "Applying removal for last keys {:?}, level {}",
                &guard_removing_poses, level
            );
            apply_removal(
                &mut cursor_guard,
                &mut guard_removing_poses,
                &mut empty_pages,
                &prev_key,
            );
        }
    }
    debug!("Have empty nodes {:?}, level {:?}", &empty_pages, level);
    empty_pages
}

pub fn level_merge<KSA, PSA, KSB, PSB>(
    src_tree: &BPlusTree<KSA, PSA>,
    dest_tree: &BPlusTree<KSB, PSB>,
) -> usize
where
    KSA: Slice<EntryKey> + Debug + 'static,
    PSA: Slice<NodeCellRef> + 'static,
    KSB: Slice<EntryKey> + Debug + 'static,
    PSB: Slice<NodeCellRef> + 'static,
{
    let left_most_leaf_guards = select::<KSA, PSA>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    debug!("Merge selected {} pages", left_most_leaf_guards.len());

    // merge to dest_tree
    {
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .cloned()
            .collect_vec();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(keys);
    }

    // cleanup upper level references
    {
        let page_keys = left_most_leaf_guards
            .iter()
            .filter(|g| !g.is_empty())
            .map(|g| g.first_key())
            .collect_vec();
        prune_selected::<KSA, PSA>(&src_tree.get_root(), page_keys, 0);
    }

    // adjust leaf left, right references
    let right_right_most = left_most_leaf_guards
        .last()
        .unwrap()
        .right_ref()
        .unwrap()
        .clone();
    let left_left_most = left_most_leaf_guards
        .first()
        .unwrap()
        .left_ref()
        .unwrap()
        .clone();
    for mut g in left_most_leaf_guards {
        *g = NodeData::Empty(box EmptyNode {
            left: Some(left_left_most.clone()),
            right: right_right_most.clone(),
        })
    }
    merge_page_len
}
<|endoftext|>


<commit_before>fn main() {
	let result = hash(515);
	for i in range(-99,99) {
		print!("{}, ", hash(i));
	}
}

//convert a number to a relatively uniform
fn hash(num: int) -> int {
	if num == 0 {
		return 0;
	}

	let mut sum = 0;
	let mut currentNum = num;

	if currentNum < 0 {
		currentNum = -currentNum;
	}

	loop {
		let digit = currentNum % 10;
		sum += digit;
		if currentNum < 10 {
			break;
		}
		currentNum /= 10;
	}

	if sum >= 10 {
		return hash(sum);
	}
	sum
}
<commit_msg>clean up part9 and add comments<commit_after>fn main() {
	for i in range(-99,99) {
		print!("{}, ", hash(i));
	}
}

//convert a number to a relatively uniform hash value from 1-9 inclusive
fn hash(num: int) -> int {
	if num == 0 {
		return 0;
	}

	let mut sum = 0;
	let mut currentNum = num;

	if currentNum < 0 {
		currentNum = -currentNum;
	}

	//sum the digits of the number
	loop {
		let digit = currentNum % 10;
		sum += digit;
		if currentNum < 10 {
			break;
		}
		currentNum /= 10;
	}

	//if there multiple digits, run the hash function again
	if sum >= 10 {
		return hash(sum);
	}
	return sum;
}
<|endoftext|>


<commit_before>use itoa;
use json;
use std::collections::BTreeMap;

pub fn canonicalize(jsn: &json::Value) -> Result<Vec<u8>, String> {
    let converted = convert(jsn)?;
    let mut buf = Vec::new();
    let _ = converted.write(&mut buf); // Vec<u8> impl always succeeds (or panics).
    Ok(buf)
}

enum Value {
    Array(Vec<Value>),
    Bool(bool),
    Null,
    Number(Number),
    Object(BTreeMap<String, Value>),
    String(String),
}

impl Value {
    fn write(&self, mut buf: &mut Vec<u8>) -> Result<(), String> {
        match self {
            &Value::Null => Ok(buf.extend(b"null")),
            &Value::Bool(true) => Ok(buf.extend(b"true")),
            &Value::Bool(false) => Ok(buf.extend(b"false")),
            &Value::Number(Number::I64(n)) => {
                itoa::write(buf, n).map(|_| ()).map_err(|err| {
                    format!("Write error: {}", err)
                })
            }
            &Value::Number(Number::U64(n)) => {
                itoa::write(buf, n).map(|_| ()).map_err(|err| {
                    format!("Write error: {}", err)
                })
            }
            &Value::String(ref s) => {
                // this mess is abusing serde_json to get json escaping
                let s = json::Value::String(s.clone());
                let s = json::to_string(&s)
                    .map_err(|e| format!("{:?}", e))?;
                Ok(buf.extend(s.as_bytes()))
            }
            &Value::Array(ref arr) => {
                buf.push(b'[');
                let mut first = true;
                for a in arr.iter() {
                    if !first {
                        buf.push(b',');
                    }
                    a.write(&mut buf)?;
                    first = false;
                }
                Ok(buf.push(b']'))
            }
            &Value::Object(ref obj) => {
                buf.push(b'{');
                let mut first = true;
                for (k, v) in obj.iter() {
                    if !first {
                        buf.push(b',');
                    }
                    first = false;

                    // this mess is abusing serde_json to get json escaping
                    let k = json::Value::String(k.clone());
                    let k = json::to_string(&k)
                        .map_err(|e| format!("{:?}", e))?;
                    buf.extend(k.as_bytes());

                    buf.push(b':');
                    v.write(&mut buf)?;
                }
                Ok(buf.push(b'}'))
            }
        }
    }
}

enum Number {
    I64(i64),
    U64(u64),
}

fn convert(jsn: &json::Value) -> Result<Value, String> {
    match jsn {
        &json::Value::Null => Ok(Value::Null),
        &json::Value::Bool(b) => Ok(Value::Bool(b)),
        &json::Value::Number(ref n) => {
            n.as_i64()
                .map(Number::I64)
                .or(n.as_u64().map(Number::U64))
                .map(Value::Number)
                .ok_or_else(|| String::from("only i64 and u64 are supported"))
        }
        &json::Value::Array(ref arr) => {
            let mut out = Vec::new();
            for res in arr.iter().map(|v| convert(v)) {
                out.push(res?)
            }
            Ok(Value::Array(out))
        }
        &json::Value::Object(ref obj) => {
            let mut out = BTreeMap::new();
            for (k, v) in obj.iter() {
                let _ = out.insert(k.clone(), convert(v)?);
            }
            Ok(Value::Object(out))
        }
        &json::Value::String(ref s) => Ok(Value::String(s.clone())),
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn write_str() {
        let jsn = Value::String(String::from("wat"));
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, b"\"wat\"");
    }

    #[test]
    fn write_arr() {
        let jsn = Value::Array(vec![
            Value::String(String::from("wat")),
            Value::String(String::from("lol")),
            Value::String(String::from("no")),
        ]);
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, b"[\"wat\",\"lol\",\"no\"]");
    }

    #[test]
    fn write_obj() {
        let mut map = BTreeMap::new();
        let arr = Value::Array(vec![
            Value::String(String::from("haha")),
            Value::String(String::from("omg so tired")),
        ]);
        let _ = map.insert(String::from("lol"), arr);
        let jsn = Value::Object(map);
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, b"{\"lol\":[\"haha\",\"omg so tired\"]}");
    }
}
<commit_msg>additional cjson test<commit_after>use itoa;
use json;
use std::collections::BTreeMap;

pub fn canonicalize(jsn: &json::Value) -> Result<Vec<u8>, String> {
    let converted = convert(jsn)?;
    let mut buf = Vec::new();
    let _ = converted.write(&mut buf); // Vec<u8> impl always succeeds (or panics).
    Ok(buf)
}

enum Value {
    Array(Vec<Value>),
    Bool(bool),
    Null,
    Number(Number),
    Object(BTreeMap<String, Value>),
    String(String),
}

impl Value {
    fn write(&self, mut buf: &mut Vec<u8>) -> Result<(), String> {
        match self {
            &Value::Null => Ok(buf.extend(b"null")),
            &Value::Bool(true) => Ok(buf.extend(b"true")),
            &Value::Bool(false) => Ok(buf.extend(b"false")),
            &Value::Number(Number::I64(n)) => {
                itoa::write(buf, n).map(|_| ()).map_err(|err| {
                    format!("Write error: {}", err)
                })
            }
            &Value::Number(Number::U64(n)) => {
                itoa::write(buf, n).map(|_| ()).map_err(|err| {
                    format!("Write error: {}", err)
                })
            }
            &Value::String(ref s) => {
                // this mess is abusing serde_json to get json escaping
                let s = json::Value::String(s.clone());
                let s = json::to_string(&s)
                    .map_err(|e| format!("{:?}", e))?;
                Ok(buf.extend(s.as_bytes()))
            }
            &Value::Array(ref arr) => {
                buf.push(b'[');
                let mut first = true;
                for a in arr.iter() {
                    if !first {
                        buf.push(b',');
                    }
                    a.write(&mut buf)?;
                    first = false;
                }
                Ok(buf.push(b']'))
            }
            &Value::Object(ref obj) => {
                buf.push(b'{');
                let mut first = true;
                for (k, v) in obj.iter() {
                    if !first {
                        buf.push(b',');
                    }
                    first = false;

                    // this mess is abusing serde_json to get json escaping
                    let k = json::Value::String(k.clone());
                    let k = json::to_string(&k)
                        .map_err(|e| format!("{:?}", e))?;
                    buf.extend(k.as_bytes());

                    buf.push(b':');
                    v.write(&mut buf)?;
                }
                Ok(buf.push(b'}'))
            }
        }
    }
}

enum Number {
    I64(i64),
    U64(u64),
}

fn convert(jsn: &json::Value) -> Result<Value, String> {
    match jsn {
        &json::Value::Null => Ok(Value::Null),
        &json::Value::Bool(b) => Ok(Value::Bool(b)),
        &json::Value::Number(ref n) => {
            n.as_i64()
                .map(Number::I64)
                .or(n.as_u64().map(Number::U64))
                .map(Value::Number)
                .ok_or_else(|| String::from("only i64 and u64 are supported"))
        }
        &json::Value::Array(ref arr) => {
            let mut out = Vec::new();
            for res in arr.iter().map(|v| convert(v)) {
                out.push(res?)
            }
            Ok(Value::Array(out))
        }
        &json::Value::Object(ref obj) => {
            let mut out = BTreeMap::new();
            for (k, v) in obj.iter() {
                let _ = out.insert(k.clone(), convert(v)?);
            }
            Ok(Value::Object(out))
        }
        &json::Value::String(ref s) => Ok(Value::String(s.clone())),
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn write_str() {
        let jsn = Value::String(String::from("wat"));
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, b"\"wat\"");
    }

    #[test]
    fn write_arr() {
        let jsn = Value::Array(vec![
            Value::String(String::from("wat")),
            Value::String(String::from("lol")),
            Value::String(String::from("no")),
        ]);
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, b"[\"wat\",\"lol\",\"no\"]");
    }

    #[test]
    fn write_obj() {
        let mut map = BTreeMap::new();
        let arr = Value::Array(vec![
            Value::String(String::from("haha")),
            Value::String(String::from("new\nline")),
        ]);
        let _ = map.insert(String::from("lol"), arr);
        let jsn = Value::Object(map);
        let mut out = Vec::new();
        jsn.write(&mut out).expect("write failed");
        assert_eq!(&out, &b"{\"lol\":[\"haha\",\"new\\nline\"]}");
    }
}
<|endoftext|>


<commit_before>//! Memory allocation APIs

#![stable(feature = "alloc_module", since = "1.28.0")]

mod global;
mod layout;

#[stable(feature = "global_alloc", since = "1.28.0")]
pub use self::global::GlobalAlloc;
#[stable(feature = "alloc_layout", since = "1.28.0")]
pub use self::layout::{Layout, LayoutErr};

use crate::fmt;
use crate::ptr::{self, NonNull};

/// The `AllocErr` error indicates an allocation failure
/// that may be due to resource exhaustion or to
/// something wrong when combining the given input arguments with this
/// allocator.
#[unstable(feature = "allocator_api", issue = "32838")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct AllocErr;

// (we need this for downstream impl of trait Error)
#[unstable(feature = "allocator_api", issue = "32838")]
impl fmt::Display for AllocErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("memory allocation failed")
    }
}

/// A desired initial state for allocated memory.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub enum AllocInit {
    /// The contents of the new memory are undefined.
    ///
    /// Reading uninitialized memory is Undefined Behavior; it must be initialized before use.
    Uninitialized,
    /// The new memory is guaranteed to be zeroed.
    Zeroed,
}

impl AllocInit {
    /// Initialize the specified memory block.
    ///
    /// This behaves like calling [`AllocInit::init_offset(memory, 0)`][off].
    ///
    /// [off]: AllocInit::init_offset
    ///
    /// # Safety
    ///
    /// * `memory.ptr` must be [valid] for writes of `memory.size` bytes.
    ///
    /// [valid]: ../../core/ptr/index.html#safety
    #[inline]
    #[unstable(feature = "allocator_api", issue = "32838")]
    pub unsafe fn init(self, memory: MemoryBlock) {
        self.init_offset(memory, 0)
    }

    /// Initialize the memory block like specified by `init` at the specified `offset`.
    ///
    /// This is a no-op for [`AllocInit::Uninitialized`][] and writes zeroes for
    /// [`AllocInit::Zeroed`][] at `ptr + offset` until `ptr + layout.size()`.
    ///
    /// # Safety
    ///
    /// * `memory.ptr` must be [valid] for writes of `memory.size` bytes.
    /// * `offset` must be smaller than or equal to `memory.size`
    ///
    /// [valid]: ../../core/ptr/index.html#safety
    #[inline]
    #[unstable(feature = "allocator_api", issue = "32838")]
    pub unsafe fn init_offset(self, memory: MemoryBlock, offset: usize) {
        debug_assert!(
            offset <= memory.size,
            "`offset` must be smaller than or equal to `memory.size`"
        );
        match self {
            AllocInit::Uninitialized => (),
            AllocInit::Zeroed => {
                memory.ptr.as_ptr().add(offset).write_bytes(0, memory.size - offset)
            }
        }
    }
}

/// Represents a block of allocated memory returned by an allocator.
#[derive(Debug, Copy, Clone)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub struct MemoryBlock {
    pub ptr: NonNull<u8>,
    pub size: usize,
}

/// A placement constraint when growing or shrinking an existing allocation.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub enum ReallocPlacement {
    /// The allocator is allowed to move the allocation to a different memory address.
    // FIXME(wg-allocators#46): Add a section to the module documentation "What is a legal
    //                          allocator" and link it at "valid location".
    ///
    /// If the allocation _does_ move, it's the responsibility of the allocator
    /// to also move the data from the previous location to the new location.
    MayMove,
    /// The address of the new memory must not change.
    ///
    /// If the allocation would have to be moved to a new location to fit, the
    /// reallocation request will fail.
    InPlace,
}

/// An implementation of `AllocRef` can allocate, grow, shrink, and deallocate arbitrary blocks of
/// data described via [`Layout`][].
///
/// `AllocRef` is designed to be implemented on ZSTs, references, or smart pointers because having
/// an allocator like `MyAlloc([u8; N])` cannot be moved, without updating the pointers to the
/// allocated memory.
///
/// Unlike [`GlobalAlloc`][], zero-sized allocations are allowed in `AllocRef`. If an underlying
/// allocator does not support this (like jemalloc) or return a null pointer (such as
/// `libc::malloc`), this case must be caught.
///
/// ### Currently allocated memory
///
/// Some of the methods require that a memory block be *currently allocated* via an allocator. This
/// means that:
///
/// * the starting address for that memory block was previously returned by [`alloc`], [`grow`], or
///   [`shrink`], and
///
/// * the memory block has not been subsequently deallocated, where blocks are either deallocated
///   directly by being passed to [`dealloc`] or were changed by being passed to [`grow`] or
///   [`shrink`] that returns `Ok`. If `grow` or `shrink` have returned `Err`, the passed pointer
///   remains valid.
///
/// [`alloc`]: AllocRef::alloc
/// [`grow`]: AllocRef::grow
/// [`shrink`]: AllocRef::shrink
/// [`dealloc`]: AllocRef::dealloc
///
/// ### Memory fitting
///
/// Some of the methods require that a layout *fit* a memory block. What it means for a layout to
/// "fit" a memory block means (or equivalently, for a memory block to "fit" a layout) is that the
/// following conditions must hold:
///
/// * The block must be allocated with the same alignment as [`layout.align()`], and
///
/// * The provided [`layout.size()`] must fall in the range `min ..= max`, where:
///   - `min` is the size of the layout most recently used to allocate the block, and
///   - `max` is the latest actual size returned from [`alloc`], [`grow`], or [`shrink`].
///
/// [`layout.align()`]: Layout::align
/// [`layout.size()`]: Layout::size
///
/// # Safety
///
/// * Memory blocks returned from an allocator must point to valid memory and retain their validity
///   until the instance and all of its clones are dropped, and
///
/// * cloning or moving the allocator must not invalidate memory blocks returned from this
///   allocator. A cloned allocator must behave like the same allocator.
///
/// * any pointer to a memory block which is [*currently allocated*] may be passed to any other
///   method of the allocator.
///
/// [*currently allocated*]: #currently-allocated-memory
#[unstable(feature = "allocator_api", issue = "32838")]
pub unsafe trait AllocRef {
    /// On success, returns a memory block meeting the size and alignment guarantees of `layout`.
    ///
    /// The returned block may have a larger size than specified by `layout.size()` and is
    /// initialized as specified by [`init`], all the way up to the returned size of the block.
    ///
    /// [`init`]: AllocInit
    ///
    /// # Errors
    ///
    /// Returning `Err` indicates that either memory is exhausted or `layout` does not meet
    /// allocator's size or alignment constraints.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr>;

    /// Deallocates the memory denoted by `memory`.
    ///
    /// # Safety
    ///
    /// `memory` must be a memory block returned by this allocator.
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout);

    /// Attempts to extend the memory block.
    ///
    /// Returns a new memory block containing a pointer and the actual size of the allocated
    /// block. The pointer is suitable for holding data described by a new layout with `layout`’s
    /// alignment and a size given by `new_size`. To accomplish this, the allocator may extend the
    /// allocation referenced by `ptr` to fit the new layout. If the [`placement`] is
    /// [`InPlace`], the returned pointer is guaranteed to be the same as the passed `ptr`.
    ///
    /// If `ReallocPlacement::MayMove` is used then ownership of the memory block referenced by `ptr`
    /// is transferred to this allocator. The memory may or may not be freed, and should be
    /// considered unusable (unless of course it is transferred back to the caller again via the
    /// return value of this method).
    ///
    /// If this method returns `Err`, then ownership of the memory block has not been transferred to
    /// this allocator, and the contents of the memory block are unaltered.
    ///
    /// The memory block will contain the following contents after a successful call to `grow`:
    ///   * Bytes `0..layout.size()` are preserved from the original allocation.
    ///   * Bytes `layout.size()..old_size` will either be preserved or initialized according to
    ///     [`init`], depending on the allocator implementation. `old_size` refers to the size of
    ///     the `MemoryBlock` prior to the `grow` call, which may be larger than the size
    ///     that was originally requested when it was allocated.
    ///   * Bytes `old_size..new_size` are initialized according to [`init`]. `new_size` refers to
    ///     the size of the `MemoryBlock` returned by the `grow` call.
    ///
    /// [`InPlace`]: ReallocPlacement::InPlace
    /// [`placement`]: ReallocPlacement
    /// [`init`]: AllocInit
    ///
    /// # Safety
    ///
    /// * `ptr` must be [*currently allocated*] via this allocator,
    /// * `layout` must [*fit*] the `ptr`. (The `new_size` argument need not fit it.)
    // We can't require that `new_size` is strictly greater than `memory.size` because of ZSTs.
    // An alternative would be
    // * `new_size must be strictly greater than `memory.size` or both are zero
    /// * `new_size` must be greater than or equal to `layout.size()`
    /// * `new_size`, when rounded up to the nearest multiple of `layout.align()`, must not overflow
    ///   (i.e., the rounded value must be less than `usize::MAX`).
    ///
    /// [*currently allocated*]: #currently-allocated-memory
    /// [*fit*]: #memory-fitting
    ///
    /// # Errors
    ///
    /// Returns `Err` if the new layout does not meet the allocator's size and alignment
    /// constraints of the allocator, or if growing otherwise fails.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove => {
                let size = layout.size();
                debug_assert!(
                    new_size >= size,
                    "`new_size` must be greater than or equal to `layout.size()`"
                );

                if new_size == size {
                    return Ok(MemoryBlock { ptr, size });
                }

                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
                let new_memory = self.alloc(new_layout, init)?;
                ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), size);
                self.dealloc(ptr, layout);
                Ok(new_memory)
            }
        }
    }

    /// Attempts to shrink the memory block.
    ///
    /// Returns a new memory block containing a pointer and the actual size of the allocated
    /// block. The pointer is suitable for holding data described by a new layout with `layout`’s
    /// alignment and a size given by `new_size`. To accomplish this, the allocator may shrink the
    /// allocation referenced by `ptr` to fit the new layout. If the [`placement`] is
    /// [`InPlace`], the returned pointer is guaranteed to be the same as the passed `ptr`.
    ///
    /// If this returns `Ok`, then ownership of the memory block referenced by `ptr` has been
    /// transferred to this allocator. The memory may or may not have been freed, and should be
    /// considered unusable unless it was transferred back to the caller again via the
    /// return value of this method.
    ///
    /// If this method returns `Err`, then ownership of the memory block has not been transferred to
    /// this allocator, and the contents of the memory block are unaltered.
    ///
    /// The behavior of how the allocator tries to shrink the memory is specified by [`placement`].
    ///
    /// [`InPlace`]: ReallocPlacement::InPlace
    /// [`placement`]: ReallocPlacement
    ///
    /// # Safety
    ///
    /// * `ptr` must be [*currently allocated*] via this allocator,
    /// * `layout` must [*fit*] the `ptr`. (The `new_size` argument need not fit it.)
    // We can't require that `new_size` is strictly smaller than `memory.size` because of ZSTs.
    // An alternative would be
    // * `new_size must be strictly smaller than `memory.size` or both are zero
    /// * `new_size` must be smaller than or equal to `layout.size()`
    ///
    /// [*currently allocated*]: #currently-allocated-memory
    /// [*fit*]: #memory-fitting
    ///
    /// # Errors
    ///
    /// Returns `Err` if the new layout does not meet the allocator's size and alignment
    /// constraints of the allocator, or if growing otherwise fails.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove => {
                let size = layout.size();
                debug_assert!(
                    new_size <= size,
                    "`new_size` must be smaller than or equal to `layout.size()`"
                );

                if new_size == size {
                    return Ok(MemoryBlock { ptr, size });
                }

                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
                let new_memory = self.alloc(new_layout, AllocInit::Uninitialized)?;
                ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), new_size);
                self.dealloc(ptr, layout);
                Ok(new_memory)
            }
        }
    }
}
<commit_msg>Rollup merge of #70691 - TimDiekmann:allocref-docs, r=RalfJung<commit_after>//! Memory allocation APIs

#![stable(feature = "alloc_module", since = "1.28.0")]

mod global;
mod layout;

#[stable(feature = "global_alloc", since = "1.28.0")]
pub use self::global::GlobalAlloc;
#[stable(feature = "alloc_layout", since = "1.28.0")]
pub use self::layout::{Layout, LayoutErr};

use crate::fmt;
use crate::ptr::{self, NonNull};

/// The `AllocErr` error indicates an allocation failure
/// that may be due to resource exhaustion or to
/// something wrong when combining the given input arguments with this
/// allocator.
#[unstable(feature = "allocator_api", issue = "32838")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct AllocErr;

// (we need this for downstream impl of trait Error)
#[unstable(feature = "allocator_api", issue = "32838")]
impl fmt::Display for AllocErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("memory allocation failed")
    }
}

/// A desired initial state for allocated memory.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub enum AllocInit {
    /// The contents of the new memory are uninitialized.
    Uninitialized,
    /// The new memory is guaranteed to be zeroed.
    Zeroed,
}

impl AllocInit {
    /// Initialize the specified memory block.
    ///
    /// This behaves like calling [`AllocInit::init_offset(memory, 0)`][off].
    ///
    /// [off]: AllocInit::init_offset
    ///
    /// # Safety
    ///
    /// * `memory.ptr` must be [valid] for writes of `memory.size` bytes.
    ///
    /// [valid]: ../../core/ptr/index.html#safety
    #[inline]
    #[unstable(feature = "allocator_api", issue = "32838")]
    pub unsafe fn init(self, memory: MemoryBlock) {
        self.init_offset(memory, 0)
    }

    /// Initialize the memory block like specified by `init` at the specified `offset`.
    ///
    /// This is a no-op for [`AllocInit::Uninitialized`][] and writes zeroes for
    /// [`AllocInit::Zeroed`][] at `ptr + offset` until `ptr + layout.size()`.
    ///
    /// # Safety
    ///
    /// * `memory.ptr` must be [valid] for writes of `memory.size` bytes.
    /// * `offset` must be smaller than or equal to `memory.size`
    ///
    /// [valid]: ../../core/ptr/index.html#safety
    #[inline]
    #[unstable(feature = "allocator_api", issue = "32838")]
    pub unsafe fn init_offset(self, memory: MemoryBlock, offset: usize) {
        debug_assert!(
            offset <= memory.size,
            "`offset` must be smaller than or equal to `memory.size`"
        );
        match self {
            AllocInit::Uninitialized => (),
            AllocInit::Zeroed => {
                memory.ptr.as_ptr().add(offset).write_bytes(0, memory.size - offset)
            }
        }
    }
}

/// Represents a block of allocated memory returned by an allocator.
#[derive(Debug, Copy, Clone)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub struct MemoryBlock {
    pub ptr: NonNull<u8>,
    pub size: usize,
}

/// A placement constraint when growing or shrinking an existing allocation.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[unstable(feature = "allocator_api", issue = "32838")]
pub enum ReallocPlacement {
    /// The allocator is allowed to move the allocation to a different memory address.
    // FIXME(wg-allocators#46): Add a section to the module documentation "What is a legal
    //                          allocator" and link it at "valid location".
    ///
    /// If the allocation _does_ move, it's the responsibility of the allocator
    /// to also move the data from the previous location to the new location.
    MayMove,
    /// The address of the new memory must not change.
    ///
    /// If the allocation would have to be moved to a new location to fit, the
    /// reallocation request will fail.
    InPlace,
}

/// An implementation of `AllocRef` can allocate, grow, shrink, and deallocate arbitrary blocks of
/// data described via [`Layout`][].
///
/// `AllocRef` is designed to be implemented on ZSTs, references, or smart pointers because having
/// an allocator like `MyAlloc([u8; N])` cannot be moved, without updating the pointers to the
/// allocated memory.
///
/// Unlike [`GlobalAlloc`][], zero-sized allocations are allowed in `AllocRef`. If an underlying
/// allocator does not support this (like jemalloc) or return a null pointer (such as
/// `libc::malloc`), this case must be caught.
///
/// ### Currently allocated memory
///
/// Some of the methods require that a memory block be *currently allocated* via an allocator. This
/// means that:
///
/// * the starting address for that memory block was previously returned by [`alloc`], [`grow`], or
///   [`shrink`], and
///
/// * the memory block has not been subsequently deallocated, where blocks are either deallocated
///   directly by being passed to [`dealloc`] or were changed by being passed to [`grow`] or
///   [`shrink`] that returns `Ok`. If `grow` or `shrink` have returned `Err`, the passed pointer
///   remains valid.
///
/// [`alloc`]: AllocRef::alloc
/// [`grow`]: AllocRef::grow
/// [`shrink`]: AllocRef::shrink
/// [`dealloc`]: AllocRef::dealloc
///
/// ### Memory fitting
///
/// Some of the methods require that a layout *fit* a memory block. What it means for a layout to
/// "fit" a memory block means (or equivalently, for a memory block to "fit" a layout) is that the
/// following conditions must hold:
///
/// * The block must be allocated with the same alignment as [`layout.align()`], and
///
/// * The provided [`layout.size()`] must fall in the range `min ..= max`, where:
///   - `min` is the size of the layout most recently used to allocate the block, and
///   - `max` is the latest actual size returned from [`alloc`], [`grow`], or [`shrink`].
///
/// [`layout.align()`]: Layout::align
/// [`layout.size()`]: Layout::size
///
/// # Safety
///
/// * Memory blocks returned from an allocator must point to valid memory and retain their validity
///   until the instance and all of its clones are dropped, and
///
/// * cloning or moving the allocator must not invalidate memory blocks returned from this
///   allocator. A cloned allocator must behave like the same allocator.
///
/// * any pointer to a memory block which is [*currently allocated*] may be passed to any other
///   method of the allocator.
///
/// [*currently allocated*]: #currently-allocated-memory
#[unstable(feature = "allocator_api", issue = "32838")]
pub unsafe trait AllocRef {
    /// On success, returns a memory block meeting the size and alignment guarantees of `layout`.
    ///
    /// The returned block may have a larger size than specified by `layout.size()` and is
    /// initialized as specified by [`init`], all the way up to the returned size of the block.
    ///
    /// [`init`]: AllocInit
    ///
    /// # Errors
    ///
    /// Returning `Err` indicates that either memory is exhausted or `layout` does not meet
    /// allocator's size or alignment constraints.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr>;

    /// Deallocates the memory denoted by `memory`.
    ///
    /// # Safety
    ///
    /// * `ptr` must be [*currently allocated*] via this allocator, and
    /// * `layout` must [*fit*] the `ptr`.
    ///
    /// [*currently allocated*]: #currently-allocated-memory
    /// [*fit*]: #memory-fitting
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout);

    /// Attempts to extend the memory block.
    ///
    /// Returns a new memory block containing a pointer and the actual size of the allocated
    /// block. The pointer is suitable for holding data described by a new layout with `layout`’s
    /// alignment and a size given by `new_size`. To accomplish this, the allocator may extend the
    /// allocation referenced by `ptr` to fit the new layout. If the [`placement`] is
    /// [`InPlace`], the returned pointer is guaranteed to be the same as the passed `ptr`.
    ///
    /// If `ReallocPlacement::MayMove` is used then ownership of the memory block referenced by `ptr`
    /// is transferred to this allocator. The memory may or may not be freed, and should be
    /// considered unusable (unless of course it is transferred back to the caller again via the
    /// return value of this method).
    ///
    /// If this method returns `Err`, then ownership of the memory block has not been transferred to
    /// this allocator, and the contents of the memory block are unaltered.
    ///
    /// The memory block will contain the following contents after a successful call to `grow`:
    ///   * Bytes `0..layout.size()` are preserved from the original allocation.
    ///   * Bytes `layout.size()..old_size` will either be preserved or initialized according to
    ///     [`init`], depending on the allocator implementation. `old_size` refers to the size of
    ///     the `MemoryBlock` prior to the `grow` call, which may be larger than the size
    ///     that was originally requested when it was allocated.
    ///   * Bytes `old_size..new_size` are initialized according to [`init`]. `new_size` refers to
    ///     the size of the `MemoryBlock` returned by the `grow` call.
    ///
    /// [`InPlace`]: ReallocPlacement::InPlace
    /// [`placement`]: ReallocPlacement
    /// [`init`]: AllocInit
    ///
    /// # Safety
    ///
    /// * `ptr` must be [*currently allocated*] via this allocator,
    /// * `layout` must [*fit*] the `ptr`. (The `new_size` argument need not fit it.)
    // We can't require that `new_size` is strictly greater than `memory.size` because of ZSTs.
    // An alternative would be
    // * `new_size must be strictly greater than `memory.size` or both are zero
    /// * `new_size` must be greater than or equal to `layout.size()`
    /// * `new_size`, when rounded up to the nearest multiple of `layout.align()`, must not overflow
    ///   (i.e., the rounded value must be less than or equal to `usize::MAX`).
    ///
    /// [*currently allocated*]: #currently-allocated-memory
    /// [*fit*]: #memory-fitting
    ///
    /// # Errors
    ///
    /// Returns `Err` if the new layout does not meet the allocator's size and alignment
    /// constraints of the allocator, or if growing otherwise fails.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove => {
                let size = layout.size();
                debug_assert!(
                    new_size >= size,
                    "`new_size` must be greater than or equal to `layout.size()`"
                );

                if new_size == size {
                    return Ok(MemoryBlock { ptr, size });
                }

                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
                let new_memory = self.alloc(new_layout, init)?;
                ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), size);
                self.dealloc(ptr, layout);
                Ok(new_memory)
            }
        }
    }

    /// Attempts to shrink the memory block.
    ///
    /// Returns a new memory block containing a pointer and the actual size of the allocated
    /// block. The pointer is suitable for holding data described by a new layout with `layout`’s
    /// alignment and a size given by `new_size`. To accomplish this, the allocator may shrink the
    /// allocation referenced by `ptr` to fit the new layout. If the [`placement`] is
    /// [`InPlace`], the returned pointer is guaranteed to be the same as the passed `ptr`.
    ///
    /// If this returns `Ok`, then ownership of the memory block referenced by `ptr` has been
    /// transferred to this allocator. The memory may or may not have been freed, and should be
    /// considered unusable unless it was transferred back to the caller again via the
    /// return value of this method.
    ///
    /// If this method returns `Err`, then ownership of the memory block has not been transferred to
    /// this allocator, and the contents of the memory block are unaltered.
    ///
    /// The behavior of how the allocator tries to shrink the memory is specified by [`placement`].
    ///
    /// [`InPlace`]: ReallocPlacement::InPlace
    /// [`placement`]: ReallocPlacement
    ///
    /// # Safety
    ///
    /// * `ptr` must be [*currently allocated*] via this allocator,
    /// * `layout` must [*fit*] the `ptr`. (The `new_size` argument need not fit it.)
    // We can't require that `new_size` is strictly smaller than `memory.size` because of ZSTs.
    // An alternative would be
    // * `new_size must be strictly smaller than `memory.size` or both are zero
    /// * `new_size` must be smaller than or equal to `layout.size()`
    ///
    /// [*currently allocated*]: #currently-allocated-memory
    /// [*fit*]: #memory-fitting
    ///
    /// # Errors
    ///
    /// Returns `Err` if the new layout does not meet the allocator's size and alignment
    /// constraints of the allocator, or if growing otherwise fails.
    ///
    /// Implementations are encouraged to return `Err` on memory exhaustion rather than panicking or
    /// aborting, but this is not a strict requirement. (Specifically: it is *legal* to implement
    /// this trait atop an underlying native allocation library that aborts on memory exhaustion.)
    ///
    /// Clients wishing to abort computation in response to an allocation error are encouraged to
    /// call the [`handle_alloc_error`] function, rather than directly invoking `panic!` or similar.
    ///
    /// [`handle_alloc_error`]: ../../alloc/alloc/fn.handle_alloc_error.html
    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        match placement {
            ReallocPlacement::InPlace => Err(AllocErr),
            ReallocPlacement::MayMove => {
                let size = layout.size();
                debug_assert!(
                    new_size <= size,
                    "`new_size` must be smaller than or equal to `layout.size()`"
                );

                if new_size == size {
                    return Ok(MemoryBlock { ptr, size });
                }

                let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
                let new_memory = self.alloc(new_layout, AllocInit::Uninitialized)?;
                ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), new_size);
                self.dealloc(ptr, layout);
                Ok(new_memory)
            }
        }
    }
}
<|endoftext|>