<commit_before>use std;
import std::io;
import std::io::{writer_util, reader_util};
import std::map;
import map::hashmap;
import result;

export rowreader, rowaccess, rowiter,
       new_reader, new_reader_readlen,
       hashmap_iter, hashmap_iter_full;

enum state {
    fieldstart(bool),
    infield(uint, uint),
    inquotedfield(uint, uint),
    inquote(uint, uint)
}

type rowreader = {
    readlen: uint,
    delim: char,
    quote: char,
    f : io::reader,
    mutable offset : uint,
    mutable buffers : [@[char]],
    mutable state : state,
    mutable trailing_nl : bool,
    mutable terminating : bool
};

type row = {
    fields : [ fieldtype ]
};

type bufferdescr = {
    escaped: bool,
    buffers: [@[char]],
    start: uint,
    end: uint
};

enum fieldtype {
    emptyfield(),
    bufferfield(bufferdescr)
}

iface rowiter {
    fn readrow() -> result::t<row, str>;
}

iface rowaccess {
    fn len() -> uint;
    fn getchars(uint) -> [char];
    fn getstr(uint) -> str;
}

fn new_reader(+f: io::reader, +delim: char, +quote: char) -> rowreader {
    {
        readlen: 1024u,
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false),
        mutable trailing_nl : false,
        mutable terminating: false
    }
}

fn new_reader_readlen(+f: io::reader, +delim: char, +quote: char, rl: uint) -> rowreader {
    {
        readlen: rl,
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false),
        mutable trailing_nl : false,
        mutable terminating: false
    }
}

impl of rowaccess for row {
    fn len() -> uint {
        vec::len(self.fields)
    }
    fn getchars(field: uint) -> [char] {
        fn unescape(escaped: [char]) -> [char] {
            let r : [char] = [];
            vec::reserve(r, vec::len(escaped));
            let in_q = false;
            for c in escaped {
                if in_q {
                    assert(c == '"');
                    in_q = false;
                } else {
                    in_q = c == '"';
                    r += [c];
                }
            }
            ret r;
        }
        alt self.fields[field] {
            emptyfield() { ret []; }
            bufferfield(desc) {
                let buf = [];
                {
                    let i = 0u;
                    while i < vec::len(desc.buffers) {
                        let from = if (i == 0u)
                            { desc.start } else { 0u };
                        let to = if (i == vec::len(desc.buffers) - 1u)
                            { desc.end } else { vec::len(*desc.buffers[i]) };
                        buf += vec::slice(*desc.buffers[i], from, to);
                        i = i + 1u;
                    }
                }
                if field == self.len() - 1u {
                    // there may be a trailing \r on the last field; we should strip it
                    // if so. bodgy here but seems the most efficient place to deal with this
                    if vec::len(buf) > 0u {
                        if buf[vec::len(buf)-1u] == '\r' {
                            buf = vec::slice(buf, 0u, vec::len(buf)-1u);
                        }
                    }
                }
                if desc.escaped {
                    buf = unescape(buf);
                }
                ret buf;
            }
        };
    }
    fn getstr(field: uint) -> str {
        ret str::from_chars(self.getchars(field));
    }
    fn getall() -> [str] {
        let a = [];
        self.map() { |s|
            a += [s];
        }
        ret a;
    }
    fn map(f: fn(s: str)) {
        let i = 0u;
        let len = self.len();
        while i < len {
            f(self.getstr(i));
            i += 1u;
        }
    }
}

impl of rowiter for rowreader {
    fn readrow() -> result::t<row, str> {
        fn statestr(state: state) -> str {
            alt state {
                fieldstart(after_delim) {
                    #fmt("fieldstart : after_delim %b", after_delim)
                }
                infield(b,o) {
                    #fmt("field : %u %u", b, o)
                }
                inquotedfield(b, o) {
                    #fmt("inquotedfield : %u %u", b, o)
                }
                inquote(b, o) {
                    #fmt("inquote : %u %u", b, o)
                }
            }
        }
        fn row_from_buf(self: rowreader, &fields: [fieldtype]) -> bool {
            fn new_bufferfield(self: rowreader, escaped: bool, sb: uint, so: uint, eo: uint) -> fieldtype {
                let eb = vec::len(self.buffers) - 1u;
                let sb = sb, so = so, eo = eo;
                //#debug("sb %u so %u eb %u eo %u", sb, so, eb, eo);
                //log(debug, vec::map(self.buffers) { |t| str::from_chars(*t) } );
                //log(debug, vec::map(self.buffers) { |t| vec::len(*t) });
                if escaped {
                    so += 1u;
                    if so > vec::len(*self.buffers[sb]) {
                        sb += 1u;
                        so = vec::len(*self.buffers[sb]) - 1u;
                    }
                    if eo > 0u {
                        eo -= 1u;
                    } else {
                        eb -= 1u;
                        eo = vec::len(*self.buffers[eb]) - 1u;
                    }
                }
                //#debug("sb %u so %u eb %u eo %u", sb, so, eb, eo);
                bufferfield({ escaped: escaped, buffers: vec::slice(self.buffers, sb, eb+1u), start: so, end: eo })
            }
            let cbuffer = vec::len(self.buffers) - 1u;
            let buf: @[char] = self.buffers[cbuffer];
            while self.offset < vec::len(*buf) {
                let coffset = self.offset;
                let c : char = buf[coffset];
                #debug("got '%c' | %s", c, statestr(self.state));
                self.offset += 1u;
                alt self.state {
                    fieldstart(after_delim) {
                        #debug("fieldstart : after_delim %b", after_delim);
                        if c == self.quote {
                            self.state = inquotedfield(cbuffer, coffset);
                        } else if c == '\n' {
                            if after_delim {
                                fields += [emptyfield];
                            }
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [emptyfield];
                        } else {
                            self.state = infield(cbuffer, coffset);
                        }
                    }
                    infield(b,o) {
                        #debug("field : %u %u", b, o);
                        if c == '\n' {
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                        }
                    }
                    inquotedfield(b, o) {
                        #debug("inquotedfield : %u %u", b, o);
                        if c == self.quote {
                            self.state = inquote(b, o);
                        }
                    }
                    inquote(b, o) {
                        #debug("inquote : %u %u", b, o);
                        if c == '\n' {
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                            ret true;
                        } else if c == self.quote {
                            self.state = inquotedfield(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                        // swallow odd chars, eg. space between field and "
                    }
                }
                #debug("now %s", statestr(self.state));
            }
            ret false;
        }

        self.state = fieldstart(false);
        let do_read = vec::len(self.buffers) == 0u;
        let fields = [];

        while !self.terminating {
            if do_read {
                let data: @[char] = @self.f.read_chars(self.readlen);
                if vec::len(*data) == 0u {
                    if !self.trailing_nl {
                        self.terminating = true;
                        data = @['\n'];
                    } else {
                        ret result::err("EOF");
                    }
                } else {
                    self.trailing_nl = data[vec::len(*data) - 1u] == '\n';
                }
                self.buffers += [data];
                self.offset = 0u;
            }

            if row_from_buf(self, fields) {
                let r: row = { fields: fields };
                fields = [];
                if vec::len(self.buffers) > 1u {
                    self.buffers = vec::slice(self.buffers, vec::len(self.buffers) - 1u, vec::len(self.buffers));
                }
                ret result::ok(r);
            }
            do_read = true;
        }
        ret result::err("unreachable");
    }
}

fn hashmap_iter_cols(r: rowreader, cols: [str], f: fn(map::hashmap<str, str>)) {
    loop {
        let res = r.readrow();
        if result::failure(res) {
            break;
        }
        let m : map::hashmap<str, str> = map::new_str_hash();
        let col = 0u;
        let row = result::get(res);
        if row.len() != vec::len(cols) {
            cont; // FIXME: how to flag that we dropped a crazy row?
        }
        result::get(res).map() { |s|
            m.insert(cols[col], s);
            col += 1u;
        };
        f(m);
    }
}

// reads the first row as a header, to derive keys for a hashmap
// emitted for each subsequent row
fn hashmap_iter(r: rowreader, f: fn(map::hashmap<str, str>)) {
    let res = r.readrow();
    alt res {
        result::ok(row) {
            hashmap_iter_cols(r, result::get(res).getall(), f);
        }
        result::err(_) { }
    }
}

// as hashmap_iter, but first apply 'hc' to each header; allows
// cleaning up headers; also allows verification that heads are
// satisfactory
fn hashmap_iter_full(r: rowreader, hmap: fn(&&h: str) -> str, hver: fn(cols: [str]) -> bool, f: fn(map::hashmap<str, str>)) {
    let res = r.readrow();
    alt res {
        result::ok(row) {
            let cols : [str] = vec::map(result::get(res).getall(), hmap);
            if !hver(cols) {
                ret;
            }
            hashmap_iter_cols(r, cols, f);
        }
        result::err(_) { }
    }
}

#[cfg(test)]
mod test {
    fn rowmatch(testdata: str, expected: [[str]]) {
        let chk = fn@(mk: fn(io::reader) -> rowreader) {
            let f = io::str_reader(testdata);
            let r = mk(f);
            let i = 0u;
            loop {
                let res = r.readrow();
                if result::failure(res) {
                    break;
                }
                let row = result::get(res);
                let expect = expected[i];

                assert(row.len() == vec::len(expect));
                let j = 0u;
                while j < row.len() {
                    assert(row.getstr(j) == expect[j]);
                    j += 1u;
                }
                i += 1u;
            }
            assert(i == vec::len(expected));
        };
        let runchecks = fn@(testdata: str) {
            // test default reader params
            chk() { |inp|
                new_reader_readlen(inp, ',', '"', 2u)
            };
            // test continuations over read buffers
            let j = 1u;
            while j < str::len(testdata) {
                chk() { |inp|
                    new_reader_readlen(inp, ',', '"', j)
                };
                j += 1u;
            }
            ret;
        };
        // so we can test trailing newline case, testdata
        // must not end in \n - leave off the last newline
        runchecks(testdata);
        runchecks(testdata+"\n");
        runchecks(str::replace(testdata, "\n", "\r\n"));
    }

    #[test]
    fn test_simple() {
        rowmatch("a,b,c,d\n1,2,3,4",
                 [["a", "b", "c", "d"], ["1", "2", "3", "4"]]);
    }

    #[test]
    fn test_trailing_comma() {
        rowmatch("a,b,c,d\n1,2,3,4,",
                 [["a", "b", "c", "d"], ["1", "2", "3", "4", ""]]);
    }

    #[test]
    fn test_leading_comma() {
        rowmatch("a,b,c,d\n,1,2,3,4",
                 [["a", "b", "c", "d"], ["", "1", "2", "3", "4"]]);
    }

    #[test]
    fn test_quote_simple() {
        rowmatch("\"Hello\",\"There\"\na,b,\"c\",d",
                 [["Hello", "There"], ["a", "b", "c", "d"]]);
    }

    #[test]
    fn test_quote_nested() {
        rowmatch("\"Hello\",\"There is a \"\"fly\"\" in my soup\"\na,b,\"c\",d",
                 [["Hello", "There is a \"fly\" in my soup"], ["a", "b", "c", "d"]]);
    }

    #[test]
    fn test_quote_with_comma() {
        rowmatch("\"1,2\"",
                 [["1,2"]])
    }

    #[test]
    fn test_quote_with_other_comma() {
        rowmatch("1,2,3,\"a,b,c\"",
                 [["1", "2", "3", "a,b,c"]])
    }

    #[test]
    fn test_blank_line() {
        rowmatch("\n\n", [[], []]);
    }
}

<commit_msg>update as io moved into core<commit_after>use std;
import io::{writer_util, reader_util};
import std::map;
import map::hashmap;
import result;

export rowreader, rowaccess, rowiter,
       new_reader, new_reader_readlen,
       hashmap_iter, hashmap_iter_full;

enum state {
    fieldstart(bool),
    infield(uint, uint),
    inquotedfield(uint, uint),
    inquote(uint, uint)
}

type rowreader = {
    readlen: uint,
    delim: char,
    quote: char,
    f : io::reader,
    mutable offset : uint,
    mutable buffers : [@[char]],
    mutable state : state,
    mutable trailing_nl : bool,
    mutable terminating : bool
};

type row = {
    fields : [ fieldtype ]
};

type bufferdescr = {
    escaped: bool,
    buffers: [@[char]],
    start: uint,
    end: uint
};

enum fieldtype {
    emptyfield(),
    bufferfield(bufferdescr)
}

iface rowiter {
    fn readrow() -> result::t<row, str>;
}

iface rowaccess {
    fn len() -> uint;
    fn getchars(uint) -> [char];
    fn getstr(uint) -> str;
}

fn new_reader(+f: io::reader, +delim: char, +quote: char) -> rowreader {
    {
        readlen: 1024u,
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false),
        mutable trailing_nl : false,
        mutable terminating: false
    }
}

fn new_reader_readlen(+f: io::reader, +delim: char, +quote: char, rl: uint) -> rowreader {
    {
        readlen: rl,
        delim: delim,
        quote: quote,
        f: f,
        mutable offset : 0u,
        mutable buffers : [],
        mutable state : fieldstart(false),
        mutable trailing_nl : false,
        mutable terminating: false
    }
}

impl of rowaccess for row {
    fn len() -> uint {
        vec::len(self.fields)
    }
    fn getchars(field: uint) -> [char] {
        fn unescape(escaped: [char]) -> [char] {
            let r : [char] = [];
            vec::reserve(r, vec::len(escaped));
            let in_q = false;
            for c in escaped {
                if in_q {
                    assert(c == '"');
                    in_q = false;
                } else {
                    in_q = c == '"';
                    r += [c];
                }
            }
            ret r;
        }
        alt self.fields[field] {
            emptyfield() { ret []; }
            bufferfield(desc) {
                let buf = [];
                {
                    let i = 0u;
                    while i < vec::len(desc.buffers) {
                        let from = if (i == 0u)
                            { desc.start } else { 0u };
                        let to = if (i == vec::len(desc.buffers) - 1u)
                            { desc.end } else { vec::len(*desc.buffers[i]) };
                        buf += vec::slice(*desc.buffers[i], from, to);
                        i = i + 1u;
                    }
                }
                if field == self.len() - 1u {
                    // there may be a trailing \r on the last field; we should strip it
                    // if so. bodgy here but seems the most efficient place to deal with this
                    if vec::len(buf) > 0u {
                        if buf[vec::len(buf)-1u] == '\r' {
                            buf = vec::slice(buf, 0u, vec::len(buf)-1u);
                        }
                    }
                }
                if desc.escaped {
                    buf = unescape(buf);
                }
                ret buf;
            }
        };
    }
    fn getstr(field: uint) -> str {
        ret str::from_chars(self.getchars(field));
    }
    fn getall() -> [str] {
        let a = [];
        self.map() { |s|
            a += [s];
        }
        ret a;
    }
    fn map(f: fn(s: str)) {
        let i = 0u;
        let len = self.len();
        while i < len {
            f(self.getstr(i));
            i += 1u;
        }
    }
}

impl of rowiter for rowreader {
    fn readrow() -> result::t<row, str> {
        fn statestr(state: state) -> str {
            alt state {
                fieldstart(after_delim) {
                    #fmt("fieldstart : after_delim %b", after_delim)
                }
                infield(b,o) {
                    #fmt("field : %u %u", b, o)
                }
                inquotedfield(b, o) {
                    #fmt("inquotedfield : %u %u", b, o)
                }
                inquote(b, o) {
                    #fmt("inquote : %u %u", b, o)
                }
            }
        }
        fn row_from_buf(self: rowreader, &fields: [fieldtype]) -> bool {
            fn new_bufferfield(self: rowreader, escaped: bool, sb: uint, so: uint, eo: uint) -> fieldtype {
                let eb = vec::len(self.buffers) - 1u;
                let sb = sb, so = so, eo = eo;
                //#debug("sb %u so %u eb %u eo %u", sb, so, eb, eo);
                //log(debug, vec::map(self.buffers) { |t| str::from_chars(*t) } );
                //log(debug, vec::map(self.buffers) { |t| vec::len(*t) });
                if escaped {
                    so += 1u;
                    if so > vec::len(*self.buffers[sb]) {
                        sb += 1u;
                        so = vec::len(*self.buffers[sb]) - 1u;
                    }
                    if eo > 0u {
                        eo -= 1u;
                    } else {
                        eb -= 1u;
                        eo = vec::len(*self.buffers[eb]) - 1u;
                    }
                }
                //#debug("sb %u so %u eb %u eo %u", sb, so, eb, eo);
                bufferfield({ escaped: escaped, buffers: vec::slice(self.buffers, sb, eb+1u), start: so, end: eo })
            }
            let cbuffer = vec::len(self.buffers) - 1u;
            let buf: @[char] = self.buffers[cbuffer];
            while self.offset < vec::len(*buf) {
                let coffset = self.offset;
                let c : char = buf[coffset];
                #debug("got '%c' | %s", c, statestr(self.state));
                self.offset += 1u;
                alt self.state {
                    fieldstart(after_delim) {
                        #debug("fieldstart : after_delim %b", after_delim);
                        if c == self.quote {
                            self.state = inquotedfield(cbuffer, coffset);
                        } else if c == '\n' {
                            if after_delim {
                                fields += [emptyfield];
                            }
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [emptyfield];
                        } else {
                            self.state = infield(cbuffer, coffset);
                        }
                    }
                    infield(b,o) {
                        #debug("field : %u %u", b, o);
                        if c == '\n' {
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                            ret true;
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, false, b, o, coffset)];
                        }
                    }
                    inquotedfield(b, o) {
                        #debug("inquotedfield : %u %u", b, o);
                        if c == self.quote {
                            self.state = inquote(b, o);
                        }
                    }
                    inquote(b, o) {
                        #debug("inquote : %u %u", b, o);
                        if c == '\n' {
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                            ret true;
                        } else if c == self.quote {
                            self.state = inquotedfield(b, o);
                        } else if c == self.delim {
                            self.state = fieldstart(true);
                            fields += [new_bufferfield(self, true, b, o, coffset)];
                        }
                        // swallow odd chars, eg. space between field and "
                    }
                }
                #debug("now %s", statestr(self.state));
            }
            ret false;
        }

        self.state = fieldstart(false);
        let do_read = vec::len(self.buffers) == 0u;
        let fields = [];

        while !self.terminating {
            if do_read {
                let data: @[char] = @self.f.read_chars(self.readlen);
                if vec::len(*data) == 0u {
                    if !self.trailing_nl {
                        self.terminating = true;
                        data = @['\n'];
                    } else {
                        ret result::err("EOF");
                    }
                } else {
                    self.trailing_nl = data[vec::len(*data) - 1u] == '\n';
                }
                self.buffers += [data];
                self.offset = 0u;
            }

            if row_from_buf(self, fields) {
                let r: row = { fields: fields };
                fields = [];
                if vec::len(self.buffers) > 1u {
                    self.buffers = vec::slice(self.buffers, vec::len(self.buffers) - 1u, vec::len(self.buffers));
                }
                ret result::ok(r);
            }
            do_read = true;
        }
        ret result::err("unreachable");
    }
}

fn hashmap_iter_cols(r: rowreader, cols: [str], f: fn(map::hashmap<str, str>)) {
    loop {
        let res = r.readrow();
        if result::failure(res) {
            break;
        }
        let m : map::hashmap<str, str> = map::new_str_hash();
        let col = 0u;
        let row = result::get(res);
        if row.len() != vec::len(cols) {
            cont; // FIXME: how to flag that we dropped a crazy row?
        }
        result::get(res).map() { |s|
            m.insert(cols[col], s);
            col += 1u;
        };
        f(m);
    }
}

// reads the first row as a header, to derive keys for a hashmap
// emitted for each subsequent row
fn hashmap_iter(r: rowreader, f: fn(map::hashmap<str, str>)) {
    let res = r.readrow();
    alt res {
        result::ok(row) {
            hashmap_iter_cols(r, result::get(res).getall(), f);
        }
        result::err(_) { }
    }
}

// as hashmap_iter, but first apply 'hc' to each header; allows
// cleaning up headers; also allows verification that heads are
// satisfactory
fn hashmap_iter_full(r: rowreader, hmap: fn(&&h: str) -> str, hver: fn(cols: [str]) -> bool, f: fn(map::hashmap<str, str>)) {
    let res = r.readrow();
    alt res {
        result::ok(row) {
            let cols : [str] = vec::map(result::get(res).getall(), hmap);
            if !hver(cols) {
                ret;
            }
            hashmap_iter_cols(r, cols, f);
        }
        result::err(_) { }
    }
}

#[cfg(test)]
mod test {
    fn rowmatch(testdata: str, expected: [[str]]) {
        let chk = fn@(mk: fn(io::reader) -> rowreader) {
            let f = io::str_reader(testdata);
            let r = mk(f);
            let i = 0u;
            loop {
                let res = r.readrow();
                if result::failure(res) {
                    break;
                }
                let row = result::get(res);
                let expect = expected[i];

                assert(row.len() == vec::len(expect));
                let j = 0u;
                while j < row.len() {
                    assert(row.getstr(j) == expect[j]);
                    j += 1u;
                }
                i += 1u;
            }
            assert(i == vec::len(expected));
        };
        let runchecks = fn@(testdata: str) {
            // test default reader params
            chk() { |inp|
                new_reader_readlen(inp, ',', '"', 2u)
            };
            // test continuations over read buffers
            let j = 1u;
            while j < str::len(testdata) {
                chk() { |inp|
                    new_reader_readlen(inp, ',', '"', j)
                };
                j += 1u;
            }
            ret;
        };
        // so we can test trailing newline case, testdata
        // must not end in \n - leave off the last newline
        runchecks(testdata);
        runchecks(testdata+"\n");
        runchecks(str::replace(testdata, "\n", "\r\n"));
    }

    #[test]
    fn test_simple() {
        rowmatch("a,b,c,d\n1,2,3,4",
                 [["a", "b", "c", "d"], ["1", "2", "3", "4"]]);
    }

    #[test]
    fn test_trailing_comma() {
        rowmatch("a,b,c,d\n1,2,3,4,",
                 [["a", "b", "c", "d"], ["1", "2", "3", "4", ""]]);
    }

    #[test]
    fn test_leading_comma() {
        rowmatch("a,b,c,d\n,1,2,3,4",
                 [["a", "b", "c", "d"], ["", "1", "2", "3", "4"]]);
    }

    #[test]
    fn test_quote_simple() {
        rowmatch("\"Hello\",\"There\"\na,b,\"c\",d",
                 [["Hello", "There"], ["a", "b", "c", "d"]]);
    }

    #[test]
    fn test_quote_nested() {
        rowmatch("\"Hello\",\"There is a \"\"fly\"\" in my soup\"\na,b,\"c\",d",
                 [["Hello", "There is a \"fly\" in my soup"], ["a", "b", "c", "d"]]);
    }

    #[test]
    fn test_quote_with_comma() {
        rowmatch("\"1,2\"",
                 [["1,2"]])
    }

    #[test]
    fn test_quote_with_other_comma() {
        rowmatch("1,2,3,\"a,b,c\"",
                 [["1", "2", "3", "a,b,c"]])
    }

    #[test]
    fn test_blank_line() {
        rowmatch("\n\n", [[], []]);
    }
}

<|endoftext|>


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

/*!
Syntax extension to generate FourCCs.

Once loaded, fourcc!() is called with a single 4-character string,
and an optional ident that is either `big`, `little`, or `target`.
The ident represents endianness, and specifies in which direction
the characters should be read. If the ident is omitted, it is assumed
to be `big`, i.e. left-to-right order. It returns a u32.

# Examples

To load the extension and use it:

```rust,ignore
#[phase(plugin)]
extern crate fourcc;

fn main() {
    let val = fourcc!("\xC0\xFF\xEE!");
    assert_eq!(val, 0xC0FFEE21u32);
    let little_val = fourcc!("foo ", little);
    assert_eq!(little_val, 0x21EEFFC0u32);
}
```

# References

* [Wikipedia: FourCC](http://en.wikipedia.org/wiki/FourCC)

*/

#![crate_id = "fourcc#0.11.0"]
#![experimental]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![license = "MIT/ASL2"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/0.11.0/")]

#![feature(plugin_registrar, managed_boxes)]

extern crate syntax;
extern crate rustc;

use syntax::ast;
use syntax::attr::contains;
use syntax::codemap::{Span, mk_sp};
use syntax::ext::base;
use syntax::ext::base::{ExtCtxt, MacExpr};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::token::InternedString;
use rustc::plugin::Registry;

use std::gc::Gc;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("fourcc", expand_syntax_ext);
}

pub fn expand_syntax_ext(cx: &mut ExtCtxt, sp: Span, tts: &[ast::TokenTree])
                         -> Box<base::MacResult> {
    let (expr, endian) = parse_tts(cx, tts);

    let little = match endian {
        None => false,
        Some(Ident{ident, span}) => match token::get_ident(ident).get() {
            "little" => true,
            "big" => false,
            "target" => target_endian_little(cx, sp),
            _ => {
                cx.span_err(span, "invalid endian directive in fourcc!");
                target_endian_little(cx, sp)
            }
        }
    };

    let s = match expr.node {
        // expression is a literal
        ast::ExprLit(ref lit) => match lit.node {
            // string literal
            ast::LitStr(ref s, _) => {
                if s.get().char_len() != 4 {
                    cx.span_err(expr.span, "string literal with len != 4 in fourcc!");
                }
                s
            }
            _ => {
                cx.span_err(expr.span, "unsupported literal in fourcc!");
                return base::DummyResult::expr(sp)
            }
        },
        _ => {
            cx.span_err(expr.span, "non-literal in fourcc!");
            return base::DummyResult::expr(sp)
        }
    };

    let mut val = 0u32;
    for codepoint in s.get().chars().take(4) {
        let byte = if codepoint as u32 > 0xFF {
            cx.span_err(expr.span, "fourcc! literal character out of range 0-255");
            0u8
        } else {
            codepoint as u8
        };

        val = if little {
            (val >> 8) | ((byte as u32) << 24)
        } else {
            (val << 8) | (byte as u32)
        };
    }
    let e = cx.expr_lit(sp, ast::LitUint(val as u64, ast::TyU32));
    MacExpr::new(e)
}

struct Ident {
    ident: ast::Ident,
    span: Span
}

fn parse_tts(cx: &ExtCtxt,
             tts: &[ast::TokenTree]) -> (Gc<ast::Expr>, Option<Ident>) {
    let p = &mut cx.new_parser_from_tts(tts);
    let ex = p.parse_expr();
    let id = if p.token == token::EOF {
        None
    } else {
        p.expect(&token::COMMA);
        let lo = p.span.lo;
        let ident = p.parse_ident();
        let hi = p.last_span.hi;
        Some(Ident{ident: ident, span: mk_sp(lo, hi)})
    };
    if p.token != token::EOF {
        p.unexpected();
    }
    (ex, id)
}

fn target_endian_little(cx: &ExtCtxt, sp: Span) -> bool {
    let meta = cx.meta_name_value(sp, InternedString::new("target_endian"),
        ast::LitStr(InternedString::new("little"), ast::CookedStr));
    contains(cx.cfg().as_slice(), meta)
}

// FIXME (10872): This is required to prevent an LLVM assert on Windows
#[test]
fn dummy_test() { }
<commit_msg>Add #[crate_name] attributes as necessary<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
Syntax extension to generate FourCCs.

Once loaded, fourcc!() is called with a single 4-character string,
and an optional ident that is either `big`, `little`, or `target`.
The ident represents endianness, and specifies in which direction
the characters should be read. If the ident is omitted, it is assumed
to be `big`, i.e. left-to-right order. It returns a u32.

# Examples

To load the extension and use it:

```rust,ignore
#[phase(plugin)]
extern crate fourcc;

fn main() {
    let val = fourcc!("\xC0\xFF\xEE!");
    assert_eq!(val, 0xC0FFEE21u32);
    let little_val = fourcc!("foo ", little);
    assert_eq!(little_val, 0x21EEFFC0u32);
}
```

# References

* [Wikipedia: FourCC](http://en.wikipedia.org/wiki/FourCC)

*/

#![crate_id = "fourcc#0.11.0"] // NOTE: remove after stage0
#![crate_name = "fourcc"]
#![experimental]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![license = "MIT/ASL2"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/0.11.0/")]
#![allow(unused_attribute)] // NOTE: remove after stage0

#![feature(plugin_registrar, managed_boxes)]

extern crate syntax;
extern crate rustc;

use syntax::ast;
use syntax::attr::contains;
use syntax::codemap::{Span, mk_sp};
use syntax::ext::base;
use syntax::ext::base::{ExtCtxt, MacExpr};
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::parse::token::InternedString;
use rustc::plugin::Registry;

use std::gc::Gc;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("fourcc", expand_syntax_ext);
}

pub fn expand_syntax_ext(cx: &mut ExtCtxt, sp: Span, tts: &[ast::TokenTree])
                         -> Box<base::MacResult> {
    let (expr, endian) = parse_tts(cx, tts);

    let little = match endian {
        None => false,
        Some(Ident{ident, span}) => match token::get_ident(ident).get() {
            "little" => true,
            "big" => false,
            "target" => target_endian_little(cx, sp),
            _ => {
                cx.span_err(span, "invalid endian directive in fourcc!");
                target_endian_little(cx, sp)
            }
        }
    };

    let s = match expr.node {
        // expression is a literal
        ast::ExprLit(ref lit) => match lit.node {
            // string literal
            ast::LitStr(ref s, _) => {
                if s.get().char_len() != 4 {
                    cx.span_err(expr.span, "string literal with len != 4 in fourcc!");
                }
                s
            }
            _ => {
                cx.span_err(expr.span, "unsupported literal in fourcc!");
                return base::DummyResult::expr(sp)
            }
        },
        _ => {
            cx.span_err(expr.span, "non-literal in fourcc!");
            return base::DummyResult::expr(sp)
        }
    };

    let mut val = 0u32;
    for codepoint in s.get().chars().take(4) {
        let byte = if codepoint as u32 > 0xFF {
            cx.span_err(expr.span, "fourcc! literal character out of range 0-255");
            0u8
        } else {
            codepoint as u8
        };

        val = if little {
            (val >> 8) | ((byte as u32) << 24)
        } else {
            (val << 8) | (byte as u32)
        };
    }
    let e = cx.expr_lit(sp, ast::LitUint(val as u64, ast::TyU32));
    MacExpr::new(e)
}

struct Ident {
    ident: ast::Ident,
    span: Span
}

fn parse_tts(cx: &ExtCtxt,
             tts: &[ast::TokenTree]) -> (Gc<ast::Expr>, Option<Ident>) {
    let p = &mut cx.new_parser_from_tts(tts);
    let ex = p.parse_expr();
    let id = if p.token == token::EOF {
        None
    } else {
        p.expect(&token::COMMA);
        let lo = p.span.lo;
        let ident = p.parse_ident();
        let hi = p.last_span.hi;
        Some(Ident{ident: ident, span: mk_sp(lo, hi)})
    };
    if p.token != token::EOF {
        p.unexpected();
    }
    (ex, id)
}

fn target_endian_little(cx: &ExtCtxt, sp: Span) -> bool {
    let meta = cx.meta_name_value(sp, InternedString::new("target_endian"),
        ast::LitStr(InternedString::new("little"), ast::CookedStr));
    contains(cx.cfg().as_slice(), meta)
}

// FIXME (10872): This is required to prevent an LLVM assert on Windows
#[test]
fn dummy_test() { }
<|endoftext|>


<commit_before>/*!
A Rust interface for Objective-C blocks.

For more information on the specifics of the block implementation, see
Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

# Invoking blocks

The `Block` struct is used for invoking blocks from Objective-C. For example,
consider this Objective-C function:

``` objc
int32_t sum(int32_t (^block)(int32_t, int32_t)) {
    return block(5, 8);
}
```

We could write it in Rust as the following:

```
fn sum(block: &mut Block<(i32, i32), i32>) {
    block.call((5, 8))
}
```

Note the extra parentheses in the `call` method, since the arguments must be
passed as a tuple.

# Creating blocks

Creating a block to pass to Objective-C can be done with the `ConcreteBlock`
struct. For example, to create a block that adds two `i32`s, we could write:

```
let block = ConcreteBlock::new(|a: i32, b: i32| a + b);
let mut block = block.copy();
assert!(block.call((5, 8)) == 13);
```

It is important to copy your block to the heap (with the `copy` method) before
passing it to Objective-C; this is because our `ConcreteBlock` is only meant
to be copied once, and we can enforce this in Rust, but if Objective-C code
were to copy it twice we could have a double free.
*/

use std::mem;
use std::ops::{Deref, DerefMut};
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::{Class, Object};
use {EncodePtr, Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &mut Block<Self, R>) -> R;
}

macro_rules! block_args_impl {
    ($($a:ident : $t:ident),*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &mut Block<Self, R>) -> R {
                let invoke: unsafe extern fn(*mut Block<Self, R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                unsafe {
                    invoke(block $(, $a)*)
                }
            }
        }
    );
}

block_args_impl!();
block_args_impl!(a: A);
block_args_impl!(a: A, b: B);
block_args_impl!(a: A, b: B, c: C);
block_args_impl!(a: A, b: B, c: C, d: D);
block_args_impl!(a: A, b: B, c: C, d: D, e: E);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A, R> where A: BlockArguments {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: unsafe extern fn(*mut Block<A, R>, ...) -> R,
}

// TODO: impl FnMut when it's possible
impl<A: BlockArguments, R> Block<A, R> where A: BlockArguments {
    /// Call self with the given arguments.
    pub fn call(&mut self, args: A) -> R {
        args.call_block(self)
    }
}

unsafe impl<A, R> Message for Block<A, R> where A: BlockArguments { }

impl<A, R> EncodePtr for Block<A, R> {
    fn ptr_code() -> &'static str { "@?" }
}

pub trait IntoConcreteBlock<A, R> where A: BlockArguments {
    fn into_concrete_block(self) -> ConcreteBlock<A, R, Self>;
}

macro_rules! concrete_block_impl {
    ($f:ident) => (
        concrete_block_impl!($f,);
    );
    ($f:ident, $($a:ident : $t:ident),*) => (
        impl<$($t,)* R, X> IntoConcreteBlock<($($t,)*), R> for X
                where X: Fn($($t,)*) -> R {
            fn into_concrete_block(self) -> ConcreteBlock<($($t,)*), R, X> {
                unsafe extern fn $f<$($t,)* R, X>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R
                        where X: Fn($($t,)*) -> R {
                    let block = &*block_ptr;
                    (block.closure)($($a),*)
                }

                unsafe {
                    ConcreteBlock::with_invoke(
                        mem::transmute($f::<$($t,)* R, X>), self)
                }
            }
        }
    );
}

concrete_block_impl!(concrete_block_invoke_args0);
concrete_block_impl!(concrete_block_invoke_args1, a: A);
concrete_block_impl!(concrete_block_invoke_args2, a: A, b: B);
concrete_block_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
concrete_block_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
concrete_block_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
concrete_block_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
concrete_block_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
concrete_block_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
concrete_block_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
concrete_block_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
concrete_block_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
concrete_block_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A, R, F> where A: BlockArguments {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, F>>>,
    closure: F,
}

impl<A, R, F> ConcreteBlock<A, R, F>
        where A: BlockArguments, F: IntoConcreteBlock<A, R> {
    /// Constructs a `ConcreteBlock` with the given closure.
    /// When the block is called, it will return the value that results from
    /// calling the closure.
    pub fn new(closure: F) -> Self {
        closure.into_concrete_block()
    }
}

impl<A, R, F> ConcreteBlock<A, R, F> where A: BlockArguments {
    /// Constructs a `ConcreteBlock` with the given invoke function and closure.
    /// Unsafe because the caller must ensure the invoke function takes the
    /// correct arguments.
    unsafe fn with_invoke(invoke: unsafe extern fn(*mut Self, ...) -> R,
            closure: F) -> Self {
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: mem::transmute(invoke),
            },
            descriptor: Box::new(BlockDescriptor::new()),
            closure: closure,
        }
    }

    /// Copy self onto the heap.
    pub fn copy(self) -> Id<Block<A, R>> {
        unsafe {
            // The copy method is declared as returning an object pointer.
            let block: *mut Object = msg_send![&self.base, copy];
            let block = block as *mut Block<A, R>;
            // At this point, our copy helper has been run so the block will
            // be moved to the heap and we can forget the original block
            // because the heap block will drop in our dispose helper.
            mem::forget(self);
            Id::from_retained_ptr(block)
        }
    }
}

impl<A, R, F> Clone for ConcreteBlock<A, R, F>
        where A: BlockArguments, F: Clone {
    fn clone(&self) -> Self {
        unsafe {
            ConcreteBlock::with_invoke(mem::transmute(self.invoke),
                self.closure.clone())
        }
    }
}

impl<A, R, F> Deref for ConcreteBlock<A, R, F>
        where A: BlockArguments {
    type Target = Block<A, R>;

    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

impl<A, R, F> DerefMut for ConcreteBlock<A, R, F>
        where A: BlockArguments {
    fn deref_mut(&mut self) -> &mut Block<A, R> {
        &mut self.base
    }
}

unsafe extern fn block_context_dispose<B>(block: &mut B) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<B>(_dst: &mut B, _src: &B) {
    // The runtime memmoves the src block into the dst block, nothing to do
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<B> BlockDescriptor<B> {
    fn new() -> BlockDescriptor<B> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<B>() as c_ulong,
            copy_helper: block_context_copy::<B>,
            dispose_helper: block_context_dispose::<B>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block_with(i: i32) -> Id<Block<(), i32>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block_with(i: i32) -> Id<Block<(i32,), i32>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &mut Block<(), i32>) -> i32 {
        let ptr = block as *mut _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *mut _)
        }
    }

    fn invoke_add_block(block: &mut Block<(i32,), i32>, a: i32) -> i32 {
        let ptr = block as *mut _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *mut _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let mut block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let mut block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        let mut block = ConcreteBlock::new(|&:| 13);
        let result = invoke_int_block(&mut block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        let mut block = ConcreteBlock::new(|&: a: i32| a + 5);
        let result = invoke_add_block(&mut block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_copy() {
        let s = "Hello!".to_string();
        let expected_len = s.len() as i32;
        let mut block = ConcreteBlock::new(move |&:| s.len() as i32);
        assert!(invoke_int_block(&mut block) == expected_len);

        let mut copied = block.copy();
        assert!(invoke_int_block(&mut copied) == expected_len);
    }
}
<commit_msg>Removed unnecessary closure annotations.<commit_after>/*!
A Rust interface for Objective-C blocks.

For more information on the specifics of the block implementation, see
Clang's documentation: http://clang.llvm.org/docs/Block-ABI-Apple.html

# Invoking blocks

The `Block` struct is used for invoking blocks from Objective-C. For example,
consider this Objective-C function:

``` objc
int32_t sum(int32_t (^block)(int32_t, int32_t)) {
    return block(5, 8);
}
```

We could write it in Rust as the following:

```
fn sum(block: &mut Block<(i32, i32), i32>) {
    block.call((5, 8))
}
```

Note the extra parentheses in the `call` method, since the arguments must be
passed as a tuple.

# Creating blocks

Creating a block to pass to Objective-C can be done with the `ConcreteBlock`
struct. For example, to create a block that adds two `i32`s, we could write:

```
let block = ConcreteBlock::new(|a: i32, b: i32| a + b);
let mut block = block.copy();
assert!(block.call((5, 8)) == 13);
```

It is important to copy your block to the heap (with the `copy` method) before
passing it to Objective-C; this is because our `ConcreteBlock` is only meant
to be copied once, and we can enforce this in Rust, but if Objective-C code
were to copy it twice we could have a double free.
*/

use std::mem;
use std::ops::{Deref, DerefMut};
use std::ptr;
use libc::{c_int, c_ulong};

use runtime::{Class, Object};
use {EncodePtr, Id, Message};

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    static _NSConcreteStackBlock: Class;
}

/// Types that may be used as the arguments to an Objective-C block.
pub trait BlockArguments {
    /// Calls the given `Block` with self as the arguments.
    fn call_block<R>(self, block: &mut Block<Self, R>) -> R;
}

macro_rules! block_args_impl {
    ($($a:ident : $t:ident),*) => (
        impl<$($t),*> BlockArguments for ($($t,)*) {
            fn call_block<R>(self, block: &mut Block<Self, R>) -> R {
                let invoke: unsafe extern fn(*mut Block<Self, R> $(, $t)*) -> R = unsafe {
                    mem::transmute(block.invoke)
                };
                let ($($a,)*) = self;
                unsafe {
                    invoke(block $(, $a)*)
                }
            }
        }
    );
}

block_args_impl!();
block_args_impl!(a: A);
block_args_impl!(a: A, b: B);
block_args_impl!(a: A, b: B, c: C);
block_args_impl!(a: A, b: B, c: C, d: D);
block_args_impl!(a: A, b: B, c: C, d: D, e: E);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
block_args_impl!(a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block that takes arguments of `A` when called and
/// returns a value of `R`.
#[repr(C)]
pub struct Block<A, R> where A: BlockArguments {
    isa: *const Class,
    flags: c_int,
    _reserved: c_int,
    invoke: unsafe extern fn(*mut Block<A, R>, ...) -> R,
}

// TODO: impl FnMut when it's possible
impl<A: BlockArguments, R> Block<A, R> where A: BlockArguments {
    /// Call self with the given arguments.
    pub fn call(&mut self, args: A) -> R {
        args.call_block(self)
    }
}

unsafe impl<A, R> Message for Block<A, R> where A: BlockArguments { }

impl<A, R> EncodePtr for Block<A, R> {
    fn ptr_code() -> &'static str { "@?" }
}

pub trait IntoConcreteBlock<A, R> where A: BlockArguments {
    fn into_concrete_block(self) -> ConcreteBlock<A, R, Self>;
}

macro_rules! concrete_block_impl {
    ($f:ident) => (
        concrete_block_impl!($f,);
    );
    ($f:ident, $($a:ident : $t:ident),*) => (
        impl<$($t,)* R, X> IntoConcreteBlock<($($t,)*), R> for X
                where X: Fn($($t,)*) -> R {
            fn into_concrete_block(self) -> ConcreteBlock<($($t,)*), R, X> {
                unsafe extern fn $f<$($t,)* R, X>(
                        block_ptr: *mut ConcreteBlock<($($t,)*), R, X>
                        $(, $a: $t)*) -> R
                        where X: Fn($($t,)*) -> R {
                    let block = &*block_ptr;
                    (block.closure)($($a),*)
                }

                unsafe {
                    ConcreteBlock::with_invoke(
                        mem::transmute($f::<$($t,)* R, X>), self)
                }
            }
        }
    );
}

concrete_block_impl!(concrete_block_invoke_args0);
concrete_block_impl!(concrete_block_invoke_args1, a: A);
concrete_block_impl!(concrete_block_invoke_args2, a: A, b: B);
concrete_block_impl!(concrete_block_invoke_args3, a: A, b: B, c: C);
concrete_block_impl!(concrete_block_invoke_args4, a: A, b: B, c: C, d: D);
concrete_block_impl!(concrete_block_invoke_args5, a: A, b: B, c: C, d: D, e: E);
concrete_block_impl!(concrete_block_invoke_args6, a: A, b: B, c: C, d: D, e: E, f: F);
concrete_block_impl!(concrete_block_invoke_args7, a: A, b: B, c: C, d: D, e: E, f: F, g: G);
concrete_block_impl!(concrete_block_invoke_args8, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H);
concrete_block_impl!(concrete_block_invoke_args9, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I);
concrete_block_impl!(concrete_block_invoke_args10, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J);
concrete_block_impl!(concrete_block_invoke_args11, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K);
concrete_block_impl!(concrete_block_invoke_args12, a: A, b: B, c: C, d: D, e: E, f: F, g: G, h: H, i: I, j: J, k: K, l: L);

/// An Objective-C block whose size is known at compile time and may be
/// constructed on the stack.
#[repr(C)]
pub struct ConcreteBlock<A, R, F> where A: BlockArguments {
    base: Block<A, R>,
    descriptor: Box<BlockDescriptor<ConcreteBlock<A, R, F>>>,
    closure: F,
}

impl<A, R, F> ConcreteBlock<A, R, F>
        where A: BlockArguments, F: IntoConcreteBlock<A, R> {
    /// Constructs a `ConcreteBlock` with the given closure.
    /// When the block is called, it will return the value that results from
    /// calling the closure.
    pub fn new(closure: F) -> Self {
        closure.into_concrete_block()
    }
}

impl<A, R, F> ConcreteBlock<A, R, F> where A: BlockArguments {
    /// Constructs a `ConcreteBlock` with the given invoke function and closure.
    /// Unsafe because the caller must ensure the invoke function takes the
    /// correct arguments.
    unsafe fn with_invoke(invoke: unsafe extern fn(*mut Self, ...) -> R,
            closure: F) -> Self {
        ConcreteBlock {
            base: Block {
                isa: &_NSConcreteStackBlock,
                // 1 << 25 = BLOCK_HAS_COPY_DISPOSE
                flags: 1 << 25,
                _reserved: 0,
                invoke: mem::transmute(invoke),
            },
            descriptor: Box::new(BlockDescriptor::new()),
            closure: closure,
        }
    }

    /// Copy self onto the heap.
    pub fn copy(self) -> Id<Block<A, R>> {
        unsafe {
            // The copy method is declared as returning an object pointer.
            let block: *mut Object = msg_send![&self.base, copy];
            let block = block as *mut Block<A, R>;
            // At this point, our copy helper has been run so the block will
            // be moved to the heap and we can forget the original block
            // because the heap block will drop in our dispose helper.
            mem::forget(self);
            Id::from_retained_ptr(block)
        }
    }
}

impl<A, R, F> Clone for ConcreteBlock<A, R, F>
        where A: BlockArguments, F: Clone {
    fn clone(&self) -> Self {
        unsafe {
            ConcreteBlock::with_invoke(mem::transmute(self.invoke),
                self.closure.clone())
        }
    }
}

impl<A, R, F> Deref for ConcreteBlock<A, R, F>
        where A: BlockArguments {
    type Target = Block<A, R>;

    fn deref(&self) -> &Block<A, R> {
        &self.base
    }
}

impl<A, R, F> DerefMut for ConcreteBlock<A, R, F>
        where A: BlockArguments {
    fn deref_mut(&mut self) -> &mut Block<A, R> {
        &mut self.base
    }
}

unsafe extern fn block_context_dispose<B>(block: &mut B) {
    // Read the block onto the stack and let it drop
    ptr::read(block);
}

unsafe extern fn block_context_copy<B>(_dst: &mut B, _src: &B) {
    // The runtime memmoves the src block into the dst block, nothing to do
}

#[repr(C)]
struct BlockDescriptor<B> {
    _reserved: c_ulong,
    block_size: c_ulong,
    copy_helper: unsafe extern fn(&mut B, &B),
    dispose_helper: unsafe extern fn(&mut B),
}

impl<B> BlockDescriptor<B> {
    fn new() -> BlockDescriptor<B> {
        BlockDescriptor {
            _reserved: 0,
            block_size: mem::size_of::<B>() as c_ulong,
            copy_helper: block_context_copy::<B>,
            dispose_helper: block_context_dispose::<B>,
        }
    }
}

#[cfg(test)]
mod tests {
    use Id;
    use objc_test_utils;
    use super::{Block, ConcreteBlock};

    fn get_int_block_with(i: i32) -> Id<Block<(), i32>> {
        unsafe {
            let ptr = objc_test_utils::get_int_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn get_add_block_with(i: i32) -> Id<Block<(i32,), i32>> {
        unsafe {
            let ptr = objc_test_utils::get_add_block_with(i);
            Id::from_retained_ptr(ptr as *mut _)
        }
    }

    fn invoke_int_block(block: &mut Block<(), i32>) -> i32 {
        let ptr = block as *mut _;
        unsafe {
            objc_test_utils::invoke_int_block(ptr as *mut _)
        }
    }

    fn invoke_add_block(block: &mut Block<(i32,), i32>, a: i32) -> i32 {
        let ptr = block as *mut _;
        unsafe {
            objc_test_utils::invoke_add_block(ptr as *mut _, a)
        }
    }

    #[test]
    fn test_call_block() {
        let mut block = get_int_block_with(13);
        assert!(block.call(()) == 13);
    }

    #[test]
    fn test_call_block_args() {
        let mut block = get_add_block_with(13);
        assert!(block.call((2,)) == 15);
    }

    #[test]
    fn test_create_block() {
        let mut block = ConcreteBlock::new(|| 13);
        let result = invoke_int_block(&mut block);
        assert!(result == 13);
    }

    #[test]
    fn test_create_block_args() {
        let mut block = ConcreteBlock::new(|a: i32| a + 5);
        let result = invoke_add_block(&mut block, 6);
        assert!(result == 11);
    }

    #[test]
    fn test_concrete_block_copy() {
        let s = "Hello!".to_string();
        let expected_len = s.len() as i32;
        let mut block = ConcreteBlock::new(move || s.len() as i32);
        assert!(invoke_int_block(&mut block) == expected_len);

        let mut copied = block.copy();
        assert!(invoke_int_block(&mut copied) == expected_len);
    }
}
<|endoftext|>


<commit_before>extern crate diffie_hellman;

use diffie_hellman::*;

#[test]
fn test_private_key_in_range_key() {
    let primes: Vec<u64> = vec![
        5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 773, 967, 3461, 6131,
    ];
    let private_keys: Vec<u64> = primes.iter().map(|x| private_key(*x)).collect();

    for i in 0..primes.len() {
        assert!(1 < private_keys[i] && private_keys[i] < primes[i]);
    }
}

#[test]
#[ignore]
fn test_public_key_correct() {
    let p: u64 = 23;
    let g: u64 = 5;

    let private_key: u64 = 6;
    let expected: u64 = 8;

    assert_eq!(public_key(p, g, private_key), expected);
}

#[test]
#[ignore]
fn test_secret_key_correct() {
    let p: u64 = 11;

    let private_key_a = 7;
    let public_key_b = 8;
    let secret = secret(p, public_key_b, private_key_a);
    let expected = 2;

    assert_eq!(secret, expected);
}

#[test]
#[ignore]
fn test_changed_secret_key() {
    let p: u64 = 13;
    let g: u64 = 11;

    let private_key_a = private_key(p);
    let private_key_b = private_key(p);

    let public_key_a = public_key(p, g, private_key_a);
    let public_key_b = public_key(p, g, private_key_b);

    // Key exchange
    let secret_a = secret(p, public_key_b, private_key_a);
    let secret_b = secret(p, public_key_a, private_key_b);

    assert_eq!(secret_a, secret_b);
}
<commit_msg>diffie-hellman: Added big numbers tests<commit_after>extern crate diffie_hellman;

use diffie_hellman::*;

#[test]
fn test_private_key_in_range_key() {
    let primes: Vec<u64> = vec![
        5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 773, 967, 3461, 6131,
    ];
    let private_keys: Vec<u64> = primes.iter().map(|x| private_key(*x)).collect();

    for i in 0..primes.len() {
        assert!(1 < private_keys[i] && private_keys[i] < primes[i]);
    }
}

#[test]
#[ignore]
fn test_public_key_correct() {
    let p: u64 = 23;
    let g: u64 = 5;

    let private_key: u64 = 6;
    let expected: u64 = 8;

    assert_eq!(public_key(p, g, private_key), expected);
}

#[test]
#[ignore]
fn test_secret_key_correct() {
    let p: u64 = 11;

    let private_key_a = 7;
    let public_key_b = 8;
    let secret = secret(p, public_key_b, private_key_a);
    let expected = 2;

    assert_eq!(secret, expected);
}

#[test]
#[ignore]
fn test_public_key_correct_big_numbers() {
    let p: u64 = 4_294_967_299;

    let g: u64 = 8;

    let private_key: u64 = 4_294_967_296;

    let expected: u64 = 4096;

    assert_eq!(public_key(p, g, private_key), expected);
}

#[test]
#[ignore]
fn test_secret_key_correct_big_numbers() {
    let p: u64 = 4_294_967_927;

    let private_key_a = 4_294_967_300;

    let public_key_b = 843;

    let secret = secret(p, public_key_b, private_key_a);

    let expected = 1_389_354_282;

    assert_eq!(secret, expected);
}

#[test]
#[ignore]
fn test_changed_secret_key() {
    let p: u64 = 13;
    let g: u64 = 11;

    let private_key_a = private_key(p);
    let private_key_b = private_key(p);

    let public_key_a = public_key(p, g, private_key_a);
    let public_key_b = public_key(p, g, private_key_b);

    // Key exchange
    let secret_a = secret(p, public_key_b, private_key_a);
    let secret_b = secret(p, public_key_a, private_key_b);

    assert_eq!(secret_a, secret_b);
}
<|endoftext|>


<commit_before>//! # Schemes
//! A scheme is a primitive for handling filesystem syscalls in Redox.
//! Schemes accept paths from the kernel for `open`, and file descriptors that they generate
//! are then passed for operations like `close`, `read`, `write`, etc.
//!
//! The kernel validates paths and file descriptors before they are passed to schemes,
//! also stripping the scheme identifier of paths if necessary.

use alloc::arc::Arc;
use alloc::boxed::Box;
use collections::BTreeMap;
use core::sync::atomic::AtomicUsize;
use spin::{Once, RwLock, RwLockReadGuard, RwLockWriteGuard};

use syscall::error::*;
use syscall::scheme::Scheme;

use self::debug::DebugScheme;
use self::event::EventScheme;
use self::env::EnvScheme;
use self::initfs::InitFsScheme;
use self::irq::IrqScheme;
use self::memory::MemoryScheme;
use self::null::NullScheme;
use self::pipe::PipeScheme;
use self::root::RootScheme;
use self::sys::SysScheme;
use self::zero::ZeroScheme;

/// `debug:` - provides access to serial console
pub mod debug;

/// `event:` - allows reading of `Event`s which are registered using `fevent`
pub mod event;

/// `env:` - access and modify environmental variables
pub mod env;

/// `initfs:` - a readonly filesystem used for initializing the system
pub mod initfs;

/// `irq:` - allows userspace handling of IRQs
pub mod irq;

/// When compiled with "live" feature - `disk:` - embedded filesystem for live disk
#[cfg(feature="live")]
pub mod live;

/// `memory:` - a scheme for accessing physical memory
pub mod memory;

/// `null:` - a scheme that will discard all writes, and read no bytes
pub mod null;

/// `pipe:` - used internally by the kernel to implement `pipe`
pub mod pipe;

/// `:` - allows the creation of userspace schemes, tightly dependent on `user`
pub mod root;

/// `sys:` - system information, such as the context list and scheme list
pub mod sys;

/// A wrapper around userspace schemes, tightly dependent on `root`
pub mod user;

/// `zero:` - a scheme that will discard all writes, and always fill read buffers with zero
pub mod zero;

/// Limit on number of schemes
pub const SCHEME_MAX_SCHEMES: usize = 65536;

/// Unique identifier for a scheme namespace.
int_like!(SchemeNamespace, AtomicSchemeNamespace, usize, AtomicUsize);

/// Unique identifier for a scheme.
int_like!(SchemeId, AtomicSchemeId, usize, AtomicUsize);

pub const ATOMIC_SCHEMEID_INIT: AtomicSchemeId = AtomicSchemeId::default();

/// Unique identifier for a file descriptor.
int_like!(FileHandle, AtomicFileHandle, usize, AtomicUsize);

/// Scheme list type
pub struct SchemeList {
    map: BTreeMap<SchemeId, Arc<Box<Scheme + Send + Sync>>>,
    names: BTreeMap<SchemeNamespace, BTreeMap<Box<[u8]>, SchemeId>>,
    next_ns: usize,
    next_id: usize
}

impl SchemeList {
    /// Create a new scheme list.
    pub fn new() -> Self {
        let mut list = SchemeList {
            map: BTreeMap::new(),
            names: BTreeMap::new(),
            next_ns: 0,
            next_id: 1
        };
        list.new_root();
        list
    }

    /// Initialize a new namespace
    fn new_ns(&mut self) -> SchemeNamespace {
        let ns = SchemeNamespace(self.next_ns);
        self.next_ns += 1;
        self.names.insert(ns, BTreeMap::new());

        self.insert(ns, Box::new(*b""), |scheme_id| Arc::new(Box::new(RootScheme::new(ns, scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"event"), |_| Arc::new(Box::new(EventScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"env"), |_| Arc::new(Box::new(EnvScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"null"), |_| Arc::new(Box::new(NullScheme))).unwrap();
        self.insert(ns, Box::new(*b"sys"), |_| Arc::new(Box::new(SysScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"zero"), |_| Arc::new(Box::new(ZeroScheme))).unwrap();

        ns
    }

    /// Initialize the root namespace
    #[cfg(not(feature="live"))]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    /// Initialize the root namespace - with live disk
    #[cfg(feature="live")]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Disk, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"disk"), |_| Arc::new(Box::new(self::live::DiskScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    pub fn make_ns(&mut self, from: SchemeNamespace, names: &[&[u8]]) -> Result<SchemeNamespace> {
        // Create an empty namespace
        let to = self.new_ns();

        // Copy requested scheme IDs
        for name in names.iter() {
            let id = if let Some((id, _scheme)) = self.get_name(from, name) {
                id
            } else {
                return Err(Error::new(ENODEV));
            };

            if let Some(ref mut names) = self.names.get_mut(&to) {
                assert!(names.insert(name.to_vec().into_boxed_slice(), id).is_none());
            } else {
                panic!("scheme namespace not found");
            }
        }

        Ok(to)
    }

    pub fn iter(&self) -> ::collections::btree_map::Iter<SchemeId, Arc<Box<Scheme + Send + Sync>>> {
        self.map.iter()
    }

    pub fn iter_name(&self, ns: SchemeNamespace) -> ::collections::btree_map::Iter<Box<[u8]>, SchemeId> {
        self.names[&ns].iter()
    }

    /// Get the nth scheme.
    pub fn get(&self, id: SchemeId) -> Option<&Arc<Box<Scheme + Send + Sync>>> {
        self.map.get(&id)
    }

    pub fn get_name(&self, ns: SchemeNamespace, name: &[u8]) -> Option<(SchemeId, &Arc<Box<Scheme + Send + Sync>>)> {
        if let Some(&id) = self.names[&ns].get(name) {
            self.get(id).map(|scheme| (id, scheme))
        } else {
            None
        }
    }

    /// Create a new scheme.
    pub fn insert<F>(&mut self, ns: SchemeNamespace, name: Box<[u8]>, scheme_fn: F) -> Result<SchemeId>
        where F: Fn(SchemeId) -> Arc<Box<Scheme + Send + Sync>>
    {
        if self.names[&ns].contains_key(&name) {
            return Err(Error::new(EEXIST));
        }

        if self.next_id >= SCHEME_MAX_SCHEMES {
            self.next_id = 1;
        }

        while self.map.contains_key(&SchemeId(self.next_id)) {
            self.next_id += 1;
        }

        /* Allow scheme list to grow if required
        if self.next_id >= SCHEME_MAX_SCHEMES {
            return Err(Error::new(EAGAIN));
        }
        */

        let id = SchemeId(self.next_id);
        self.next_id += 1;

        let scheme = scheme_fn(id);

        assert!(self.map.insert(id, scheme).is_none());
        if let Some(ref mut names) = self.names.get_mut(&ns) {
            assert!(names.insert(name, id).is_none());
        } else {
            panic!("scheme namespace not found");
        }
        Ok(id)
    }
}

/// Schemes list
static SCHEMES: Once<RwLock<SchemeList>> = Once::new();

/// Initialize schemes, called if needed
fn init_schemes() -> RwLock<SchemeList> {
    RwLock::new(SchemeList::new())
}

/// Get the global schemes list, const
pub fn schemes() -> RwLockReadGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).read()
}

/// Get the global schemes list, mutable
pub fn schemes_mut() -> RwLockWriteGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).write()
}
<commit_msg>Add memory scheme to live configuration<commit_after>//! # Schemes
//! A scheme is a primitive for handling filesystem syscalls in Redox.
//! Schemes accept paths from the kernel for `open`, and file descriptors that they generate
//! are then passed for operations like `close`, `read`, `write`, etc.
//!
//! The kernel validates paths and file descriptors before they are passed to schemes,
//! also stripping the scheme identifier of paths if necessary.

use alloc::arc::Arc;
use alloc::boxed::Box;
use collections::BTreeMap;
use core::sync::atomic::AtomicUsize;
use spin::{Once, RwLock, RwLockReadGuard, RwLockWriteGuard};

use syscall::error::*;
use syscall::scheme::Scheme;

use self::debug::DebugScheme;
use self::event::EventScheme;
use self::env::EnvScheme;
use self::initfs::InitFsScheme;
use self::irq::IrqScheme;
use self::memory::MemoryScheme;
use self::null::NullScheme;
use self::pipe::PipeScheme;
use self::root::RootScheme;
use self::sys::SysScheme;
use self::zero::ZeroScheme;

/// `debug:` - provides access to serial console
pub mod debug;

/// `event:` - allows reading of `Event`s which are registered using `fevent`
pub mod event;

/// `env:` - access and modify environmental variables
pub mod env;

/// `initfs:` - a readonly filesystem used for initializing the system
pub mod initfs;

/// `irq:` - allows userspace handling of IRQs
pub mod irq;

/// When compiled with "live" feature - `disk:` - embedded filesystem for live disk
#[cfg(feature="live")]
pub mod live;

/// `memory:` - a scheme for accessing physical memory
pub mod memory;

/// `null:` - a scheme that will discard all writes, and read no bytes
pub mod null;

/// `pipe:` - used internally by the kernel to implement `pipe`
pub mod pipe;

/// `:` - allows the creation of userspace schemes, tightly dependent on `user`
pub mod root;

/// `sys:` - system information, such as the context list and scheme list
pub mod sys;

/// A wrapper around userspace schemes, tightly dependent on `root`
pub mod user;

/// `zero:` - a scheme that will discard all writes, and always fill read buffers with zero
pub mod zero;

/// Limit on number of schemes
pub const SCHEME_MAX_SCHEMES: usize = 65536;

/// Unique identifier for a scheme namespace.
int_like!(SchemeNamespace, AtomicSchemeNamespace, usize, AtomicUsize);

/// Unique identifier for a scheme.
int_like!(SchemeId, AtomicSchemeId, usize, AtomicUsize);

pub const ATOMIC_SCHEMEID_INIT: AtomicSchemeId = AtomicSchemeId::default();

/// Unique identifier for a file descriptor.
int_like!(FileHandle, AtomicFileHandle, usize, AtomicUsize);

/// Scheme list type
pub struct SchemeList {
    map: BTreeMap<SchemeId, Arc<Box<Scheme + Send + Sync>>>,
    names: BTreeMap<SchemeNamespace, BTreeMap<Box<[u8]>, SchemeId>>,
    next_ns: usize,
    next_id: usize
}

impl SchemeList {
    /// Create a new scheme list.
    pub fn new() -> Self {
        let mut list = SchemeList {
            map: BTreeMap::new(),
            names: BTreeMap::new(),
            next_ns: 0,
            next_id: 1
        };
        list.new_root();
        list
    }

    /// Initialize a new namespace
    fn new_ns(&mut self) -> SchemeNamespace {
        let ns = SchemeNamespace(self.next_ns);
        self.next_ns += 1;
        self.names.insert(ns, BTreeMap::new());

        self.insert(ns, Box::new(*b""), |scheme_id| Arc::new(Box::new(RootScheme::new(ns, scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"event"), |_| Arc::new(Box::new(EventScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"env"), |_| Arc::new(Box::new(EnvScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"null"), |_| Arc::new(Box::new(NullScheme))).unwrap();
        self.insert(ns, Box::new(*b"sys"), |_| Arc::new(Box::new(SysScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"zero"), |_| Arc::new(Box::new(ZeroScheme))).unwrap();

        ns
    }

    /// Initialize the root namespace
    #[cfg(not(feature="live"))]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    /// Initialize the root namespace - with live disk
    #[cfg(feature="live")]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Disk, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"disk"), |_| Arc::new(Box::new(self::live::DiskScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    pub fn make_ns(&mut self, from: SchemeNamespace, names: &[&[u8]]) -> Result<SchemeNamespace> {
        // Create an empty namespace
        let to = self.new_ns();

        // Copy requested scheme IDs
        for name in names.iter() {
            let id = if let Some((id, _scheme)) = self.get_name(from, name) {
                id
            } else {
                return Err(Error::new(ENODEV));
            };

            if let Some(ref mut names) = self.names.get_mut(&to) {
                assert!(names.insert(name.to_vec().into_boxed_slice(), id).is_none());
            } else {
                panic!("scheme namespace not found");
            }
        }

        Ok(to)
    }

    pub fn iter(&self) -> ::collections::btree_map::Iter<SchemeId, Arc<Box<Scheme + Send + Sync>>> {
        self.map.iter()
    }

    pub fn iter_name(&self, ns: SchemeNamespace) -> ::collections::btree_map::Iter<Box<[u8]>, SchemeId> {
        self.names[&ns].iter()
    }

    /// Get the nth scheme.
    pub fn get(&self, id: SchemeId) -> Option<&Arc<Box<Scheme + Send + Sync>>> {
        self.map.get(&id)
    }

    pub fn get_name(&self, ns: SchemeNamespace, name: &[u8]) -> Option<(SchemeId, &Arc<Box<Scheme + Send + Sync>>)> {
        if let Some(&id) = self.names[&ns].get(name) {
            self.get(id).map(|scheme| (id, scheme))
        } else {
            None
        }
    }

    /// Create a new scheme.
    pub fn insert<F>(&mut self, ns: SchemeNamespace, name: Box<[u8]>, scheme_fn: F) -> Result<SchemeId>
        where F: Fn(SchemeId) -> Arc<Box<Scheme + Send + Sync>>
    {
        if self.names[&ns].contains_key(&name) {
            return Err(Error::new(EEXIST));
        }

        if self.next_id >= SCHEME_MAX_SCHEMES {
            self.next_id = 1;
        }

        while self.map.contains_key(&SchemeId(self.next_id)) {
            self.next_id += 1;
        }

        /* Allow scheme list to grow if required
        if self.next_id >= SCHEME_MAX_SCHEMES {
            return Err(Error::new(EAGAIN));
        }
        */

        let id = SchemeId(self.next_id);
        self.next_id += 1;

        let scheme = scheme_fn(id);

        assert!(self.map.insert(id, scheme).is_none());
        if let Some(ref mut names) = self.names.get_mut(&ns) {
            assert!(names.insert(name, id).is_none());
        } else {
            panic!("scheme namespace not found");
        }
        Ok(id)
    }
}

/// Schemes list
static SCHEMES: Once<RwLock<SchemeList>> = Once::new();

/// Initialize schemes, called if needed
fn init_schemes() -> RwLock<SchemeList> {
    RwLock::new(SchemeList::new())
}

/// Get the global schemes list, const
pub fn schemes() -> RwLockReadGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).read()
}

/// Get the global schemes list, mutable
pub fn schemes_mut() -> RwLockWriteGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).write()
}
<|endoftext|>