<commit_before>#[doc = "Rust TNetStrings serialization library."];

use std;
import option::{some, none};
import std::map;
import std::map::hashmap;
import std::rand;

import io::writer_util;
import io::reader_util;

export t;
export str;
export int;
export float;
export bool;
export null;
export map;
export vec;
export to_writer;
export to_bytes;
export to_str;
export from_reader;
export from_bytes;
export from_str;

#[doc = "Represents a tnetstring value."]
enum t {
    #[doc = "str"]
    str([u8]),
    #[doc = "int"]
    int(int),
    #[doc = "floating"]
    float(float),
    #[doc = "boolean"]
    bool(bool),
    #[doc = "null"]
    null,
    #[doc = "map"]
    map(map::hashmap<[u8], t>),
    #[doc = "list"]
    vec([t]),
}

#[doc = "Serializes a tnetstring value into a io::writer."]
fn to_writer(writer: io::writer, t: t) {
    alt t {
        str(s) {
            writer.write_str(#fmt("%u:", vec::len(s)));
            writer.write(s);
            writer.write_char(',');
        }
        int(i) {
            let s = int::str(i);
            writer.write_str(#fmt("%u:%s#", str::len(s), s));
        }
        float(f) {
            let s = float::to_str(f, 6u);
            writer.write_str(#fmt("%u:%s^", str::len(s), s));
        }
        bool(b) {
            let s = bool::to_str(b);
            writer.write_str(#fmt("%u:%s!", str::len(s), s));
        }
        map(m) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            m.items({ |key, value|
                to_writer(wr, str(key));
                to_writer(wr, value);
            });
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", vec::len(payload)));
            writer.write(payload);
            writer.write_char('}');
        }
        vec(l) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            vec::iter(l, {|e| to_writer(wr, e) });
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", vec::len(payload)));
            writer.write(payload);
            writer.write_char(']');
        }
        null {
            writer.write_str("0:~");
        }
    }
}

#[doc = "Serializes a tnetstring value into a byte string."]
fn to_bytes(t: t) -> [u8] {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, t);
    io::mem_buffer_buf(buf)
}

#[doc = "Serializes a tnetstring value into a string."]
fn to_str(t: t) -> str {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, t);
    io::mem_buffer_str(buf)
}

#[doc = "Deserializes a tnetstring value from an io::reader."]
fn from_reader(reader: io::reader) -> option<t> {
    assert !reader.eof();

    let mut c = reader.read_byte();
    let mut len = 0u;

    // Note that netstring spec explicitly forbids padding zeros.
    // If the first char is zero, it must be the only char.
    if c < '0' as int || c > '9' as int {
        fail "Not a tnetstring: invalid or missing length prefix";
    } else if c == '0' as int {
        c = reader.read_byte();
    } else {
        do {
            len = (10u * len) + ((c as uint) - ('0' as uint));

            if reader.eof() {
                fail "Not a tnetstring: invalid or missing length prefix";
            }
            c = reader.read_byte();
        } while c >= '0' as int && c <= '9' as int;
    }

    // Validate end-of-length-prefix marker.
    if c != ':' as int {
        fail "Not a tnetstring: missing length prefix";
    }

    // Read the data plus terminating type tag.
    let payload = reader.read_bytes(len);

    if vec::len(payload) != len {
        fail "Not a tnetstring: invalid length prefix";
    }

    if reader.eof() {
        fail "Not a tnetstring: missing type tag";
    }

    alt reader.read_byte() as char {
      '#' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(int::from_str(s)) {|v| some(int(v)) }
      }
      '}' { some(map(parse_map(payload))) }
      ']' { some(vec(parse_vec(payload))) }
      '!' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(bool::from_str(s)) {|v| some(bool(v)) }
      }
      '^' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(float::from_str(s)) {|v| some(float(v)) }
      }
      '~' {
        assert vec::len(payload) == 0u;
        some(null)
      }
      ',' { some(str(payload)) }
      c {
        let s = str::from_char(c);
        fail #fmt("Invalid payload type: %?", s)
      }
    }
}

fn parse_vec(data: [u8]) -> [t] {
    if vec::len(data) == 0u { ret []; }

    let reader = io::bytes_reader(data);

    let value = from_reader(reader);
    assert option::is_some(value);
    let mut result = [option::get(value)];

    while !reader.eof() {
        let value = from_reader(reader);
        assert option::is_some(value);
        vec::push(result, option::get(value));
    }

    ret result;
}

fn parse_pair(reader: io::reader) -> ([u8], t) {
    let key = from_reader(reader);
    assert option::is_some(key);
    assert !reader.eof();

    let key = alt option::get(key) {
        str(key) { key }
        _ { fail "Keys can only be strings." }
    };

    let value = from_reader(reader);
    assert option::is_some(value);

    ret (key, option::get(value));
}

fn parse_map(data: [u8]) -> map::hashmap<[u8], t> {
    if vec::len(data) == 0u { ret map::bytes_hash(); }

    let reader = io::bytes_reader(data);
    let (key, value) = parse_pair(reader);

    let result = map::bytes_hash();
    result.insert(key, value);

    while !reader.eof() {
        let (key, value) = parse_pair(reader);
        result.insert(key, value);
    }

    ret result;
}

#[doc = "Deserializes a tnetstring value from a byte string."]
fn from_bytes(data: [u8]) -> (option<t>, [u8]) {
    let rdr = io::bytes_reader(data);
    let tnetstring = from_reader(rdr);
    (tnetstring, rdr.read_whole_stream())
}

#[doc = "Deserializes a tnetstring value from a string."]
fn from_str(data: str) -> (option<t>, str) {
    io::with_str_reader(data) {|rdr|
        let tnetstring = from_reader(rdr);
        let bytes = rdr.read_whole_stream();
        (tnetstring, str::from_bytes(bytes))
    }

}

#[doc = "Test the equality between two tnetstring values"]
fn eq(t0: t, t1: t) -> bool {
    alt (t0, t1) {
        (str(s0), str(s1)) { s0 == s1 }
        (int(i0), int(i1)) { i0 == i1 }
        (float(f0), float(f1)) { f0 == f1 }
        (bool(b0), bool(b1)) { b0 == b1 }
        (null, null) { true }
        (map(d0), map(d1)) {
            if d0.size() == d1.size() {
                let mut equal = true;
                d0.items() { |k, v|
                    if !d1.contains_key(k) || !eq(d1.get(k), v) {
                        equal = false;
                    }
                };
                equal
            } else {
                false
            }
        }
        (vec(l0), vec(l1)) {
            vec::all2(l0, l1, eq)
        }
        _ { false }
    }
}

#[cfg(test)]
mod tests {
    // Tests inspired by https://github.com/rfk/tnetstring.

    fn test(string: str, expected: t) {
        let (actual, rest) = from_str(string);
        assert option::is_some(actual);
        assert rest == "";

        let actual = option::get(actual);
        assert eq(actual, expected);
        assert to_str(expected) == string;
    }

    #[test]
    fn test_format() {
        test("11:hello world,", str(str::bytes("hello world")));
        test("0:}", map(map::bytes_hash()));
        test("0:]", vec([]));

        let d = map::bytes_hash();
        d.insert(str::bytes("hello"),
                vec([
                    int(12345678901),
                    str(str::bytes("this")),
                    bool(true),
                    null,
                    str(str::bytes("\x00\x00\x00\x00"))]));

        test("51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00" +
            "\x00,]}", map(d));

        test("5:12345#", int(12345));
        test("12:this is cool,", str(str::bytes("this is cool")));
        test("0:,", str(str::bytes("")));
        test("0:~", null);
        test("4:true!", bool(true));
        test("5:false!", bool(false));
        test("10:\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00,",
            str(str::bytes("\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")));
        test("24:5:12345#5:67890#5:xxxxx,]",
            vec(
                [int(12345),
                int(67890),
                str(str::bytes("xxxxx"))]));
        test("18:3:0.1^3:0.2^3:0.4^]",
           vec([float(0.1), float(0.2), float(0.4)]));
        test("243:238:233:228:223:218:213:208:203:198:193:188:183:178:173:" +
            "168:163:158:153:148:143:138:133:128:123:118:113:108:103:99:95:" +
            "91:87:83:79:75:71:67:63:59:55:51:47:43:39:35:31:27:23:19:15:" +
            "11:hello-there,]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]" +
            "]]]]",
            vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([
                    str(str::bytes("hello-there"))
                ])])])])])])])])])])])])])])])])])])])])])])])])])])])])
                ])])])])])])])])])])])])])])])])])])])])])])]));
    }

    #[test]
    fn test_random() {
        fn randint(rng: rand::rng, a: u32, b: u32) -> u32 {
            (rng.next() % (b - a + 1u32)) + a
        }

        fn get_random_object(rng: rand::rng, depth: u32) -> t {
            if randint(rng, depth, 10u32) <= 4u32 {
                if randint(rng, 0u32, 1u32) == 0u32 {
                    let n = randint(rng, 0u32, 10u32);
                    vec(vec::from_fn(n as uint) { |_i|
                        get_random_object(rng, depth + 1u32)
                    })
                } else {
                    let d = map::bytes_hash();

                    let mut i = randint(rng, 0u32, 10u32);
                    while i != 0u32 {
                        let s = rng.gen_bytes(randint(rng, 0u32, 100u32) as uint);
                        d.insert(s, get_random_object(rng, depth + 1u32));
                        i -= 1u32;
                    }
                    map(d)
                }
            } else {
                alt randint(rng, 0u32, 5u32) {
                  0u32 { null }
                  1u32 { bool(true) }
                  2u32 { bool(false) }
                  3u32 {
                    if randint(rng, 0u32, 1u32) == 0u32 {
                        int(rng.next() as int)
                    } else {
                        int(-rng.next() as int)
                    }
                  }
                  4u32 {
                    let mut f = rng.next_float();

                    // Generate a float that can be exactly converted to
                    // and from a string.
                    while true {
                        alt float::from_str(float::to_str(f, 6u)) {
                          some(f1) {
                            if f == f1 { break; }
                            f = f1;
                          }
                          none { fail }
                        }
                    }

                    if randint(rng, 0u32, 1u32) == 0u32 {
                        float(f)
                    } else {
                        float(-f)
                    }
                  }
                  5u32 {
                    str(rng.gen_bytes(randint(rng, 0u32, 100u32) as uint))
                  }
                  _ { fail }
                }
            }
        }

        let rng = rand::rng();

        let mut i = 500u;
        while i != 0u {
            let v0 = get_random_object(rng, 0u32);

            alt from_bytes(to_bytes(v0)) {
                (some(v1), rest) if rest == [] {
                    assert eq(v0, v1);
                }
                _ { fail; }
            }
            i -= 1u;
        }
    }
}
<commit_msg>Switch to for-style loops<commit_after>#[doc = "Rust TNetStrings serialization library."];

use std;
import option::{some, none};
import std::map;
import std::map::hashmap;
import std::rand;

import io::writer_util;
import io::reader_util;

export t;
export str;
export int;
export float;
export bool;
export null;
export map;
export vec;
export to_writer;
export to_bytes;
export to_str;
export from_reader;
export from_bytes;
export from_str;

#[doc = "Represents a tnetstring value."]
enum t {
    #[doc = "str"]
    str([u8]),
    #[doc = "int"]
    int(int),
    #[doc = "floating"]
    float(float),
    #[doc = "boolean"]
    bool(bool),
    #[doc = "null"]
    null,
    #[doc = "map"]
    map(map::hashmap<[u8], t>),
    #[doc = "list"]
    vec([t]),
}

#[doc = "Serializes a tnetstring value into a io::writer."]
fn to_writer(writer: io::writer, t: t) {
    alt t {
        str(s) {
            writer.write_str(#fmt("%u:", vec::len(s)));
            writer.write(s);
            writer.write_char(',');
        }
        int(i) {
            let s = int::str(i);
            writer.write_str(#fmt("%u:%s#", str::len(s), s));
        }
        float(f) {
            let s = float::to_str(f, 6u);
            writer.write_str(#fmt("%u:%s^", str::len(s), s));
        }
        bool(b) {
            let s = bool::to_str(b);
            writer.write_str(#fmt("%u:%s!", str::len(s), s));
        }
        map(m) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            for m.each { |key, value|
                to_writer(wr, str(key));
                to_writer(wr, value);
            }
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", vec::len(payload)));
            writer.write(payload);
            writer.write_char('}');
        }
        vec(l) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            vec::iter(l, {|e| to_writer(wr, e) });
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", vec::len(payload)));
            writer.write(payload);
            writer.write_char(']');
        }
        null {
            writer.write_str("0:~");
        }
    }
}

#[doc = "Serializes a tnetstring value into a byte string."]
fn to_bytes(t: t) -> [u8] {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, t);
    io::mem_buffer_buf(buf)
}

#[doc = "Serializes a tnetstring value into a string."]
fn to_str(t: t) -> str {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, t);
    io::mem_buffer_str(buf)
}

#[doc = "Deserializes a tnetstring value from an io::reader."]
fn from_reader(reader: io::reader) -> option<t> {
    assert !reader.eof();

    let mut c = reader.read_byte();
    let mut len = 0u;

    // Note that netstring spec explicitly forbids padding zeros.
    // If the first char is zero, it must be the only char.
    if c < '0' as int || c > '9' as int {
        fail "Not a tnetstring: invalid or missing length prefix";
    } else if c == '0' as int {
        c = reader.read_byte();
    } else {
        do {
            len = (10u * len) + ((c as uint) - ('0' as uint));

            if reader.eof() {
                fail "Not a tnetstring: invalid or missing length prefix";
            }
            c = reader.read_byte();
        } while c >= '0' as int && c <= '9' as int;
    }

    // Validate end-of-length-prefix marker.
    if c != ':' as int {
        fail "Not a tnetstring: missing length prefix";
    }

    // Read the data plus terminating type tag.
    let payload = reader.read_bytes(len);

    if vec::len(payload) != len {
        fail "Not a tnetstring: invalid length prefix";
    }

    if reader.eof() {
        fail "Not a tnetstring: missing type tag";
    }

    alt reader.read_byte() as char {
      '#' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(int::from_str(s)) {|v| some(int(v)) }
      }
      '}' { some(map(parse_map(payload))) }
      ']' { some(vec(parse_vec(payload))) }
      '!' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(bool::from_str(s)) {|v| some(bool(v)) }
      }
      '^' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(float::from_str(s)) {|v| some(float(v)) }
      }
      '~' {
        assert vec::len(payload) == 0u;
        some(null)
      }
      ',' { some(str(payload)) }
      c {
        let s = str::from_char(c);
        fail #fmt("Invalid payload type: %?", s)
      }
    }
}

fn parse_vec(data: [u8]) -> [t] {
    if vec::len(data) == 0u { ret []; }

    let reader = io::bytes_reader(data);

    let value = from_reader(reader);
    assert option::is_some(value);
    let mut result = [option::get(value)];

    while !reader.eof() {
        let value = from_reader(reader);
        assert option::is_some(value);
        vec::push(result, option::get(value));
    }

    ret result;
}

fn parse_pair(reader: io::reader) -> ([u8], t) {
    let key = from_reader(reader);
    assert option::is_some(key);
    assert !reader.eof();

    let key = alt option::get(key) {
        str(key) { key }
        _ { fail "Keys can only be strings." }
    };

    let value = from_reader(reader);
    assert option::is_some(value);

    ret (key, option::get(value));
}

fn parse_map(data: [u8]) -> map::hashmap<[u8], t> {
    if vec::len(data) == 0u { ret map::bytes_hash(); }

    let reader = io::bytes_reader(data);
    let (key, value) = parse_pair(reader);

    let result = map::bytes_hash();
    result.insert(key, value);

    while !reader.eof() {
        let (key, value) = parse_pair(reader);
        result.insert(key, value);
    }

    ret result;
}

#[doc = "Deserializes a tnetstring value from a byte string."]
fn from_bytes(data: [u8]) -> (option<t>, [u8]) {
    let rdr = io::bytes_reader(data);
    let tnetstring = from_reader(rdr);
    (tnetstring, rdr.read_whole_stream())
}

#[doc = "Deserializes a tnetstring value from a string."]
fn from_str(data: str) -> (option<t>, str) {
    io::with_str_reader(data) {|rdr|
        let tnetstring = from_reader(rdr);
        let bytes = rdr.read_whole_stream();
        (tnetstring, str::from_bytes(bytes))
    }

}

#[doc = "Test the equality between two tnetstring values"]
fn eq(t0: t, t1: t) -> bool {
    alt (t0, t1) {
        (str(s0), str(s1)) { s0 == s1 }
        (int(i0), int(i1)) { i0 == i1 }
        (float(f0), float(f1)) { f0 == f1 }
        (bool(b0), bool(b1)) { b0 == b1 }
        (null, null) { true }
        (map(d0), map(d1)) {
          if d0.size() == d1.size() {
              for d0.each() { |k, v|
                  if !d1.contains_key(k) || !eq(d1.get(k), v) {
                      ret false;
                  }
              }
              true
          } else {
              false
          }
        }
        (vec(l0), vec(l1)) { vec::all2(l0, l1, eq) }
        _ { false }
    }
}

#[cfg(test)]
mod tests {
    // Tests inspired by https://github.com/rfk/tnetstring.

    fn test(string: str, expected: t) {
        let (actual, rest) = from_str(string);
        assert option::is_some(actual);
        assert rest == "";

        let actual = option::get(actual);
        assert eq(actual, expected);
        assert to_str(expected) == string;
    }

    #[test]
    fn test_format() {
        test("11:hello world,", str(str::bytes("hello world")));
        test("0:}", map(map::bytes_hash()));
        test("0:]", vec([]));

        let d = map::bytes_hash();
        d.insert(str::bytes("hello"),
                vec([
                    int(12345678901),
                    str(str::bytes("this")),
                    bool(true),
                    null,
                    str(str::bytes("\x00\x00\x00\x00"))]));

        test("51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00" +
            "\x00,]}", map(d));

        test("5:12345#", int(12345));
        test("12:this is cool,", str(str::bytes("this is cool")));
        test("0:,", str(str::bytes("")));
        test("0:~", null);
        test("4:true!", bool(true));
        test("5:false!", bool(false));
        test("10:\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00,",
            str(str::bytes("\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")));
        test("24:5:12345#5:67890#5:xxxxx,]",
            vec(
                [int(12345),
                int(67890),
                str(str::bytes("xxxxx"))]));
        test("18:3:0.1^3:0.2^3:0.4^]",
           vec([float(0.1), float(0.2), float(0.4)]));
        test("243:238:233:228:223:218:213:208:203:198:193:188:183:178:173:" +
            "168:163:158:153:148:143:138:133:128:123:118:113:108:103:99:95:" +
            "91:87:83:79:75:71:67:63:59:55:51:47:43:39:35:31:27:23:19:15:" +
            "11:hello-there,]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]" +
            "]]]]",
            vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([vec([vec([vec([vec([vec([vec(
                [vec([vec([
                    str(str::bytes("hello-there"))
                ])])])])])])])])])])])])])])])])])])])])])])])])])])])])
                ])])])])])])])])])])])])])])])])])])])])])])]));
    }

    #[test]
    fn test_random() {
        fn randint(rng: rand::rng, a: u32, b: u32) -> u32 {
            (rng.next() % (b - a + 1u32)) + a
        }

        fn get_random_object(rng: rand::rng, depth: u32) -> t {
            if randint(rng, depth, 10u32) <= 4u32 {
                if randint(rng, 0u32, 1u32) == 0u32 {
                    let n = randint(rng, 0u32, 10u32);
                    vec(vec::from_fn(n as uint) { |_i|
                        get_random_object(rng, depth + 1u32)
                    })
                } else {
                    let d = map::bytes_hash();

                    let mut i = randint(rng, 0u32, 10u32);
                    while i != 0u32 {
                        let s = rng.gen_bytes(randint(rng, 0u32, 100u32) as uint);
                        d.insert(s, get_random_object(rng, depth + 1u32));
                        i -= 1u32;
                    }
                    map(d)
                }
            } else {
                alt randint(rng, 0u32, 5u32) {
                  0u32 { null }
                  1u32 { bool(true) }
                  2u32 { bool(false) }
                  3u32 {
                    if randint(rng, 0u32, 1u32) == 0u32 {
                        int(rng.next() as int)
                    } else {
                        int(-rng.next() as int)
                    }
                  }
                  4u32 {
                    let mut f = rng.next_float();

                    // Generate a float that can be exactly converted to
                    // and from a string.
                    while true {
                        alt float::from_str(float::to_str(f, 6u)) {
                          some(f1) {
                            if f == f1 { break; }
                            f = f1;
                          }
                          none { fail }
                        }
                    }

                    if randint(rng, 0u32, 1u32) == 0u32 {
                        float(f)
                    } else {
                        float(-f)
                    }
                  }
                  5u32 {
                    str(rng.gen_bytes(randint(rng, 0u32, 100u32) as uint))
                  }
                  _ { fail }
                }
            }
        }

        let rng = rand::rng();

        let mut i = 500u;
        while i != 0u {
            let v0 = get_random_object(rng, 0u32);

            alt from_bytes(to_bytes(v0)) {
                (some(v1), rest) if rest == [] {
                    assert eq(v0, v1);
                }
                _ { fail; }
            }
            i -= 1u;
        }
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::ops::Deref;
use std::rc::{Rc, Weak};

use super::super::ir;
use super::super::target::bytecode;

// Manages loading IR units and compiling them into callable functions.
//
// NOTE: Eventually this should probably handle loading source files. Really
//   the whole continuum between source and specialized native code.
pub struct Loader {
    units: Vec<LoadedUnit>,
}

impl Loader {
    pub fn new() -> Self {
        Self { units: vec![] }
    }

    pub fn load(&mut self, ir_unit: ir::layout::Unit) -> LoadedUnit {
        let loaded_unit = LoadedUnit::empty();
        let functions = ir::compiler::compile(&ir_unit)
            .functions
            .into_iter()
            .map(|function| {
                LoadedFunction(Rc::new(InnerLoadedFunction {
                    unit: Rc::downgrade(&loaded_unit.0),
                    function,
                }))
            })
            .collect::<Vec<LoadedFunction>>();
        loaded_unit.0.borrow_mut().functions = functions;

        self.units.push(loaded_unit.clone());
        loaded_unit
    }
}

// Opaque wrapper around a reference-counted loaded unit.
#[derive(Clone)]
pub struct LoadedUnit(Rc<RefCell<InnerLoadedUnit>>);

type WeakLoadedUnit = Weak<RefCell<InnerLoadedUnit>>;

impl LoadedUnit {
    fn empty() -> Self {
        InnerLoadedUnit { functions: vec![] }.into()
    }

    pub fn main(&self) -> LoadedFunction {
        self.0.borrow().functions[0].clone()
    }

    pub fn function(&self, id: u16) -> LoadedFunction {
        let this = &self.0;
        this.borrow()
            .functions
            .iter()
            .find(|&function| function.id() == id)
            .expect("Function not found")
            .clone()
    }
}

impl From<InnerLoadedUnit> for LoadedUnit {
    fn from(loaded_unit: InnerLoadedUnit) -> LoadedUnit {
        LoadedUnit(Rc::new(RefCell::new(loaded_unit)))
    }
}

pub struct InnerLoadedUnit {
    functions: Vec<LoadedFunction>,
}

// Handle to a loaded bytecode function and its unit. Used by `Frame`.
#[derive(Clone)]
pub struct InnerLoadedFunction {
    unit: WeakLoadedUnit,
    function: bytecode::layout::Function,
}

impl InnerLoadedFunction {
    #[inline]
    pub fn id(&self) -> u16 {
        self.function.id
    }

    #[inline]
    pub fn registers(&self) -> u8 {
        self.function.registers
    }

    #[inline]
    pub fn locals(&self) -> u8 {
        self.function.locals
    }

    #[inline]
    pub fn instruction(
        &self,
        basic_block_address: usize,
        instruction_address: usize,
    ) -> bytecode::layout::Instruction {
        let block = &self.function.basic_blocks[basic_block_address];
        block.instructions[instruction_address].clone()
    }

    pub fn locals_names(&self) -> Vec<String> {
        self.function.locals_names.clone()
    }

    pub fn unit(&self) -> LoadedUnit {
        let upgraded = self.unit.upgrade().expect("Unit has been dropped");
        LoadedUnit(upgraded)
    }
}

#[derive(Clone)]
pub struct LoadedFunction(Rc<InnerLoadedFunction>);

impl Deref for LoadedFunction {
    type Target = InnerLoadedFunction;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<commit_msg>Do conversions via `.into()`<commit_after>use std::cell::RefCell;
use std::ops::Deref;
use std::rc::{Rc, Weak};

use super::super::ir;
use super::super::target::bytecode;

// Manages loading IR units and compiling them into callable functions.
//
// NOTE: Eventually this should probably handle loading source files. Really
//   the whole continuum between source and specialized native code.
pub struct Loader {
    units: Vec<LoadedUnit>,
}

impl Loader {
    pub fn new() -> Self {
        Self { units: vec![] }
    }

    pub fn load(&mut self, ir_unit: ir::layout::Unit) -> LoadedUnit {
        let loaded_unit = LoadedUnit::empty();
        let functions = ir::compiler::compile(&ir_unit)
            .functions
            .into_iter()
            .map(|function| {
                InnerLoadedFunction {
                    unit: Rc::downgrade(&loaded_unit.0),
                    function,
                }
                .into()
            })
            .collect::<Vec<LoadedFunction>>();
        loaded_unit.0.borrow_mut().functions = functions;

        self.units.push(loaded_unit.clone());
        loaded_unit
    }
}

// Opaque wrapper around a reference-counted loaded unit.
#[derive(Clone)]
pub struct LoadedUnit(Rc<RefCell<InnerLoadedUnit>>);

type WeakLoadedUnit = Weak<RefCell<InnerLoadedUnit>>;

impl LoadedUnit {
    fn empty() -> Self {
        InnerLoadedUnit { functions: vec![] }.into()
    }

    pub fn main(&self) -> LoadedFunction {
        self.0.borrow().functions[0].clone()
    }

    pub fn function(&self, id: u16) -> LoadedFunction {
        let this = &self.0;
        this.borrow()
            .functions
            .iter()
            .find(|&function| function.id() == id)
            .expect("Function not found")
            .clone()
    }
}

impl From<InnerLoadedUnit> for LoadedUnit {
    fn from(loaded_unit: InnerLoadedUnit) -> LoadedUnit {
        LoadedUnit(Rc::new(RefCell::new(loaded_unit)))
    }
}

pub struct InnerLoadedUnit {
    functions: Vec<LoadedFunction>,
}

// Handle to a loaded bytecode function and its unit. Used by `Frame`.
#[derive(Clone)]
pub struct InnerLoadedFunction {
    unit: WeakLoadedUnit,
    function: bytecode::layout::Function,
}

impl InnerLoadedFunction {
    #[inline]
    pub fn id(&self) -> u16 {
        self.function.id
    }

    #[inline]
    pub fn registers(&self) -> u8 {
        self.function.registers
    }

    #[inline]
    pub fn locals(&self) -> u8 {
        self.function.locals
    }

    #[inline]
    pub fn instruction(
        &self,
        basic_block_address: usize,
        instruction_address: usize,
    ) -> bytecode::layout::Instruction {
        let block = &self.function.basic_blocks[basic_block_address];
        block.instructions[instruction_address].clone()
    }

    pub fn locals_names(&self) -> Vec<String> {
        self.function.locals_names.clone()
    }

    pub fn unit(&self) -> LoadedUnit {
        let upgraded = self.unit.upgrade().expect("Unit has been dropped");
        LoadedUnit(upgraded)
    }
}

#[derive(Clone)]
pub struct LoadedFunction(Rc<InnerLoadedFunction>);

impl From<InnerLoadedFunction> for LoadedFunction {
    fn from(loaded_function: InnerLoadedFunction) -> LoadedFunction {
        LoadedFunction(Rc::new(loaded_function))
    }
}

impl Deref for LoadedFunction {
    type Target = InnerLoadedFunction;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<|endoftext|>


<commit_before>//! Unsafe operations

export reinterpret_cast, forget, bump_box_refcount, transmute;
export transmute_mut, transmute_immut, transmute_region, transmute_mut_region;

export SharedMutableState, shared_mutable_state, clone_shared_mutable_state;
export get_shared_mutable_state, get_shared_immutable_state;
export Exclusive, exclusive;

import task::atomically;

#[abi = "rust-intrinsic"]
extern mod rusti {
    fn forget<T>(-x: T);
    fn reinterpret_cast<T, U>(e: T) -> U;
}

/// Casts the value at `src` to U. The two types must have the same length.
#[inline(always)]
unsafe fn reinterpret_cast<T, U>(src: T) -> U {
    rusti::reinterpret_cast(src)
}

/**
 * Move a thing into the void
 *
 * The forget function will take ownership of the provided value but neglect
 * to run any required cleanup or memory-management operations on it. This
 * can be used for various acts of magick, particularly when using
 * reinterpret_cast on managed pointer types.
 */
#[inline(always)]
unsafe fn forget<T>(-thing: T) { rusti::forget(thing); }

/**
 * Force-increment the reference count on a shared box. If used
 * uncarefully, this can leak the box. Use this in conjunction with transmute
 * and/or reinterpret_cast when such calls would otherwise scramble a box's
 * reference count
 */
unsafe fn bump_box_refcount<T>(+t: @T) { forget(t); }

/**
 * Transform a value of one type into a value of another type.
 * Both types must have the same size and alignment.
 *
 * # Example
 *
 *     assert transmute("L") == ~[76u8, 0u8];
 */
unsafe fn transmute<L, G>(-thing: L) -> G {
    let newthing = reinterpret_cast(thing);
    forget(thing);
    return newthing;
}

/// Coerce an immutable reference to be mutable.
unsafe fn transmute_mut<T>(+ptr: &T) -> &mut T { transmute(ptr) }
/// Coerce a mutable reference to be immutable.
unsafe fn transmute_immut<T>(+ptr: &mut T) -> &T { transmute(ptr) }
/// Coerce a borrowed pointer to have an arbitrary associated region.
unsafe fn transmute_region<T>(+ptr: &a/T) -> &b/T { transmute(ptr) }
/// Coerce a borrowed mutable pointer to have an arbitrary associated region.
unsafe fn transmute_mut_region<T>(+ptr: &a/mut T) -> &b/mut T {
    transmute(ptr)
}

/****************************************************************************
 * Shared state & exclusive ARC
 ****************************************************************************/

type ArcData<T> = {
    mut count: libc::intptr_t,
    data: T
};

struct ArcDestruct<T> {
   let data: *libc::c_void;
   new(data: *libc::c_void) { self.data = data; }
   drop unsafe {
      let data: ~ArcData<T> = unsafe::reinterpret_cast(self.data);
      let new_count = rustrt::rust_atomic_decrement(&mut data.count);
      assert new_count >= 0;
      if new_count == 0 {
          // drop glue takes over.
      } else {
        unsafe::forget(data);
      }
   }
}

/**
 * COMPLETELY UNSAFE. Used as a primitive for the safe versions in std::arc.
 *
 * Data races between tasks can result in crashes and, with sufficient
 * cleverness, arbitrary type coercion.
 */
type SharedMutableState<T: send> = ArcDestruct<T>;

unsafe fn shared_mutable_state<T: send>(+data: T) -> SharedMutableState<T> {
    let data = ~{mut count: 1, data: data};
    unsafe {
        let ptr = unsafe::transmute(data);
        ArcDestruct(ptr)
    }
}

#[inline(always)]
unsafe fn get_shared_mutable_state<T: send>(rc: &SharedMutableState<T>)
        -> &mut T {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        assert ptr.count > 0;
        // Cast us back into the correct region
        let r = unsafe::reinterpret_cast(&ptr.data);
        unsafe::forget(ptr);
        return r;
    }
}
#[inline(always)]
unsafe fn get_shared_immutable_state<T: send>(rc: &SharedMutableState<T>)
        -> &T {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        assert ptr.count > 0;
        // Cast us back into the correct region
        let r = unsafe::reinterpret_cast(&ptr.data);
        unsafe::forget(ptr);
        return r;
    }
}

unsafe fn clone_shared_mutable_state<T: send>(rc: &SharedMutableState<T>)
        -> SharedMutableState<T> {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
        assert new_count >= 2;
        unsafe::forget(ptr);
    }
    ArcDestruct((*rc).data)
}

/****************************************************************************/

#[allow(non_camel_case_types)] // runtime type
type rust_little_lock = *libc::c_void;

#[abi = "cdecl"]
extern mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    fn rust_create_little_lock() -> rust_little_lock;
    fn rust_destroy_little_lock(lock: rust_little_lock);
    fn rust_lock_little_lock(lock: rust_little_lock);
    fn rust_unlock_little_lock(lock: rust_little_lock);
}

struct LittleLock {
    let l: rust_little_lock;
    new() {
        self.l = rustrt::rust_create_little_lock();
    }
    drop { rustrt::rust_destroy_little_lock(self.l); }
}

impl LittleLock {
    #[inline(always)]
    unsafe fn lock<T>(f: fn() -> T) -> T {
        struct Unlock {
            let l: rust_little_lock;
            new(l: rust_little_lock) { self.l = l; }
            drop { rustrt::rust_unlock_little_lock(self.l); }
        }

        do atomically {
            rustrt::rust_lock_little_lock(self.l);
            let _r = Unlock(self.l);
            f()
        }
    }
}

struct ExData<T: send> { lock: LittleLock; mut failed: bool; mut data: T; }
/**
 * An arc over mutable data that is protected by a lock. For library use only.
 */
struct Exclusive<T: send> { x: SharedMutableState<ExData<T>>; }

fn exclusive<T:send >(+user_data: T) -> Exclusive<T> {
    let data = ExData {
        lock: LittleLock(), mut failed: false, mut data: user_data
    };
    Exclusive { x: unsafe { shared_mutable_state(data) } }
}

impl<T: send> Exclusive<T> {
    // Duplicate an exclusive ARC, as std::arc::clone.
    fn clone() -> Exclusive<T> {
        Exclusive { x: unsafe { clone_shared_mutable_state(&self.x) } }
    }

    // Exactly like std::arc::mutex_arc,access(), but with the little_lock
    // instead of a proper mutex. Same reason for being unsafe.
    //
    // Currently, scheduling operations (i.e., yielding, receiving on a pipe,
    // accessing the provided condition variable) are prohibited while inside
    // the exclusive. Supporting that is a work in progress.
    #[inline(always)]
    unsafe fn with<U>(f: fn(x: &mut T) -> U) -> U {
        let rec = unsafe { get_shared_mutable_state(&self.x) };
        do rec.lock.lock {
            if rec.failed {
                fail ~"Poisoned exclusive - another task failed inside!";
            }
            rec.failed = true;
            let result = f(&mut rec.data);
            rec.failed = false;
            result
        }
    }
}

/****************************************************************************
 * Tests
 ****************************************************************************/

#[cfg(test)]
mod tests {

    #[test]
    fn test_reinterpret_cast() {
        assert unsafe { reinterpret_cast(1) } == 1u;
    }

    #[test]
    fn test_bump_box_refcount() {
        unsafe {
            let box = @~"box box box";       // refcount 1
            bump_box_refcount(box);         // refcount 2
            let ptr: *int = transmute(box); // refcount 2
            let _box1: @~str = reinterpret_cast(ptr);
            let _box2: @~str = reinterpret_cast(ptr);
            assert *_box1 == ~"box box box";
            assert *_box2 == ~"box box box";
            // Will destroy _box1 and _box2. Without the bump, this would
            // use-after-free. With too many bumps, it would leak.
        }
    }

    #[test]
    fn test_transmute() {
        unsafe {
            let x = @1;
            let x: *int = transmute(x);
            assert *x == 1;
            let _x: @int = transmute(x);
        }
    }

    #[test]
    fn test_transmute2() {
        unsafe {
            assert transmute(~"L") == ~[76u8, 0u8];
        }
    }

    #[test]
    fn exclusive_arc() {
        let mut futures = ~[];

        let num_tasks = 10u;
        let count = 10u;

        let total = exclusive(~mut 0u);

        for uint::range(0u, num_tasks) |_i| {
            let total = total.clone();
            futures += ~[future::spawn(|| {
                for uint::range(0u, count) |_i| {
                    do total.with |count| {
                        **count += 1u;
                    }
                }
            })];
        };

        for futures.each |f| { f.get() }

        do total.with |total| {
            assert **total == num_tasks * count
        };
    }

    #[test] #[should_fail] #[ignore(cfg(windows))]
    fn exclusive_poison() {
        // Tests that if one task fails inside of an exclusive, subsequent
        // accesses will also fail.
        let x = exclusive(1);
        let x2 = x.clone();
        do task::try {
            do x2.with |one| {
                assert *one == 2;
            }
        };
        do x.with |one| {
            assert *one == 1;
        }
    }
}
<commit_msg>implement arc::unwrap and add exclusive_unwrap tests<commit_after>//! Unsafe operations

export reinterpret_cast, forget, bump_box_refcount, transmute;
export transmute_mut, transmute_immut, transmute_region, transmute_mut_region;

export SharedMutableState, shared_mutable_state, clone_shared_mutable_state;
export get_shared_mutable_state, get_shared_immutable_state;
export unwrap_shared_mutable_state;
export Exclusive, exclusive;

import task::atomically;

#[abi = "rust-intrinsic"]
extern mod rusti {
    fn forget<T>(-x: T);
    fn reinterpret_cast<T, U>(e: T) -> U;
}

/// Casts the value at `src` to U. The two types must have the same length.
#[inline(always)]
unsafe fn reinterpret_cast<T, U>(src: T) -> U {
    rusti::reinterpret_cast(src)
}

/**
 * Move a thing into the void
 *
 * The forget function will take ownership of the provided value but neglect
 * to run any required cleanup or memory-management operations on it. This
 * can be used for various acts of magick, particularly when using
 * reinterpret_cast on managed pointer types.
 */
#[inline(always)]
unsafe fn forget<T>(-thing: T) { rusti::forget(thing); }

/**
 * Force-increment the reference count on a shared box. If used
 * uncarefully, this can leak the box. Use this in conjunction with transmute
 * and/or reinterpret_cast when such calls would otherwise scramble a box's
 * reference count
 */
unsafe fn bump_box_refcount<T>(+t: @T) { forget(t); }

/**
 * Transform a value of one type into a value of another type.
 * Both types must have the same size and alignment.
 *
 * # Example
 *
 *     assert transmute("L") == ~[76u8, 0u8];
 */
unsafe fn transmute<L, G>(-thing: L) -> G {
    let newthing = reinterpret_cast(thing);
    forget(thing);
    return newthing;
}

/// Coerce an immutable reference to be mutable.
unsafe fn transmute_mut<T>(+ptr: &T) -> &mut T { transmute(ptr) }
/// Coerce a mutable reference to be immutable.
unsafe fn transmute_immut<T>(+ptr: &mut T) -> &T { transmute(ptr) }
/// Coerce a borrowed pointer to have an arbitrary associated region.
unsafe fn transmute_region<T>(+ptr: &a/T) -> &b/T { transmute(ptr) }
/// Coerce a borrowed mutable pointer to have an arbitrary associated region.
unsafe fn transmute_mut_region<T>(+ptr: &a/mut T) -> &b/mut T {
    transmute(ptr)
}

/****************************************************************************
 * Shared state & exclusive ARC
 ****************************************************************************/

// An unwrapper uses this protocol to communicate with the "other" task that
// drops the last refcount on an arc. Unfortunately this can't be a proper
// pipe protocol because the unwrapper has to access both stages at once.
type UnwrapProto = ~mut option<(pipes::chan_one<()>, pipes::port_one<bool>)>;

struct ArcData<T> {
    mut count:     libc::intptr_t;
    mut unwrapper: libc::uintptr_t; // either a UnwrapProto or 0
    // FIXME(#3224) should be able to make this non-option to save memory, and
    // in unwrap() use "let ~ArcData { data: result, _ } = thing" to unwrap it
    mut data:      option<T>;
}

struct ArcDestruct<T> {
    mut data: *libc::c_void;
    new(data: *libc::c_void) { self.data = data; }
    drop unsafe {
        if self.data.is_null() {
            return; // Happens when destructing an unwrapper's handle.
        }
        do task::unkillable {
            let data: ~ArcData<T> = unsafe::reinterpret_cast(self.data);
            let new_count = rustrt::rust_atomic_decrement(&mut data.count);
            assert new_count >= 0;
            if new_count == 0 {
                // Were we really last, or should we hand off to an unwrapper?
                // It's safe to not xchg because the unwrapper will set the
                // unwrap lock *before* dropping his/her reference. In effect,
                // being here means we're the only *awake* task with the data.
                if data.unwrapper != 0 {
                    let p: UnwrapProto =
                        unsafe::reinterpret_cast(data.unwrapper);
                    let (message, response) = option::swap_unwrap(p);
                    // Send 'ready' and wait for a response.
                    pipes::send_one(message, ());
                    // Unkillable wait. Message guaranteed to come.
                    if pipes::recv_one(response) {
                        // Other task got the data.
                        unsafe::forget(data);
                    } else {
                        // Other task was killed. drop glue takes over.
                    }
                } else {
                    // drop glue takes over.
                }
            } else {
                unsafe::forget(data);
            }
        }
    }
}

unsafe fn unwrap_shared_mutable_state<T: send>(+rc: SharedMutableState<T>)
        -> T {
    struct DeathThroes<T> {
        mut ptr:      option<~ArcData<T>>;
        mut response: option<pipes::chan_one<bool>>;
        drop unsafe {
            let response = option::swap_unwrap(&mut self.response);
            // In case we get killed early, we need to tell the person who
            // tried to wake us whether they should hand-off the data to us.
            if task::failing() {
                pipes::send_one(response, false);
                // Either this swap_unwrap or the one below (at "Got here")
                // ought to run.
                unsafe::forget(option::swap_unwrap(&mut self.ptr));
            } else {
                assert self.ptr.is_none();
                pipes::send_one(response, true);
            }
        }
    }

    do task::unkillable {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast(rc.data);
        let (c1,p1) = pipes::oneshot(); // ()
        let (c2,p2) = pipes::oneshot(); // bool
        let server: UnwrapProto = ~mut some((c1,p2));
        let serverp: libc::uintptr_t = unsafe::transmute(server);
        // Try to put our server end in the unwrapper slot.
        if rustrt::rust_compare_and_swap_ptr(&mut ptr.unwrapper, 0, serverp) {
            // Got in. Step 0: Tell destructor not to run. We are now it.
            rc.data = ptr::null();
            // Step 1 - drop our own reference.
            let new_count = rustrt::rust_atomic_decrement(&mut ptr.count);
            assert new_count >= 0;
            if new_count == 0 {
                // We were the last owner. Can unwrap immediately.
                // Also we have to free the server endpoints.
                let _server: UnwrapProto = unsafe::transmute(serverp);
                option::swap_unwrap(&mut ptr.data)
                // drop glue takes over.
            } else {
                // The *next* person who sees the refcount hit 0 will wake us.
                let end_result =
                    DeathThroes { ptr: some(ptr), response: some(c2) };
                let mut p1 = some(p1); // argh
                do task::rekillable {
                    pipes::recv_one(option::swap_unwrap(&mut p1));
                }
                // Got here. Back in the 'unkillable' without getting killed.
                // Recover ownership of ptr, then take the data out.
                let ptr = option::swap_unwrap(&mut end_result.ptr);
                option::swap_unwrap(&mut ptr.data)
                // drop glue takes over.
            }
        } else {
            // Somebody else was trying to unwrap. Avoid guaranteed deadlock.
            unsafe::forget(ptr);
            // Also we have to free the (rejected) server endpoints.
            let _server: UnwrapProto = unsafe::transmute(serverp);
            fail ~"Another task is already unwrapping this ARC!";
        }
    }
}

/**
 * COMPLETELY UNSAFE. Used as a primitive for the safe versions in std::arc.
 *
 * Data races between tasks can result in crashes and, with sufficient
 * cleverness, arbitrary type coercion.
 */
type SharedMutableState<T: send> = ArcDestruct<T>;

unsafe fn shared_mutable_state<T: send>(+data: T) -> SharedMutableState<T> {
    let data = ~ArcData { count: 1, unwrapper: 0, data: some(data) };
    unsafe {
        let ptr = unsafe::transmute(data);
        ArcDestruct(ptr)
    }
}

#[inline(always)]
unsafe fn get_shared_mutable_state<T: send>(rc: &SharedMutableState<T>)
        -> &mut T {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        assert ptr.count > 0;
        // Cast us back into the correct region
        let r = unsafe::transmute_region(option::get_ref(&ptr.data));
        unsafe::forget(ptr);
        return unsafe::transmute_mut(r);
    }
}
#[inline(always)]
unsafe fn get_shared_immutable_state<T: send>(rc: &SharedMutableState<T>)
        -> &T {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        assert ptr.count > 0;
        // Cast us back into the correct region
        let r = unsafe::transmute_region(option::get_ref(&ptr.data));
        unsafe::forget(ptr);
        return r;
    }
}

unsafe fn clone_shared_mutable_state<T: send>(rc: &SharedMutableState<T>)
        -> SharedMutableState<T> {
    unsafe {
        let ptr: ~ArcData<T> = unsafe::reinterpret_cast((*rc).data);
        let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
        assert new_count >= 2;
        unsafe::forget(ptr);
    }
    ArcDestruct((*rc).data)
}

/****************************************************************************/

#[allow(non_camel_case_types)] // runtime type
type rust_little_lock = *libc::c_void;

#[abi = "cdecl"]
extern mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_compare_and_swap_ptr(address: &mut libc::uintptr_t,
                                 oldval: libc::uintptr_t,
                                 newval: libc::uintptr_t) -> bool;

    fn rust_create_little_lock() -> rust_little_lock;
    fn rust_destroy_little_lock(lock: rust_little_lock);
    fn rust_lock_little_lock(lock: rust_little_lock);
    fn rust_unlock_little_lock(lock: rust_little_lock);
}

struct LittleLock {
    let l: rust_little_lock;
    new() {
        self.l = rustrt::rust_create_little_lock();
    }
    drop { rustrt::rust_destroy_little_lock(self.l); }
}

impl LittleLock {
    #[inline(always)]
    unsafe fn lock<T>(f: fn() -> T) -> T {
        struct Unlock {
            let l: rust_little_lock;
            new(l: rust_little_lock) { self.l = l; }
            drop { rustrt::rust_unlock_little_lock(self.l); }
        }

        do atomically {
            rustrt::rust_lock_little_lock(self.l);
            let _r = Unlock(self.l);
            f()
        }
    }
}

struct ExData<T: send> { lock: LittleLock; mut failed: bool; mut data: T; }
/**
 * An arc over mutable data that is protected by a lock. For library use only.
 */
struct Exclusive<T: send> { x: SharedMutableState<ExData<T>>; }

fn exclusive<T:send >(+user_data: T) -> Exclusive<T> {
    let data = ExData {
        lock: LittleLock(), mut failed: false, mut data: user_data
    };
    Exclusive { x: unsafe { shared_mutable_state(data) } }
}

impl<T: send> Exclusive<T> {
    // Duplicate an exclusive ARC, as std::arc::clone.
    fn clone() -> Exclusive<T> {
        Exclusive { x: unsafe { clone_shared_mutable_state(&self.x) } }
    }

    // Exactly like std::arc::mutex_arc,access(), but with the little_lock
    // instead of a proper mutex. Same reason for being unsafe.
    //
    // Currently, scheduling operations (i.e., yielding, receiving on a pipe,
    // accessing the provided condition variable) are prohibited while inside
    // the exclusive. Supporting that is a work in progress.
    #[inline(always)]
    unsafe fn with<U>(f: fn(x: &mut T) -> U) -> U {
        let rec = unsafe { get_shared_mutable_state(&self.x) };
        do rec.lock.lock {
            if rec.failed {
                fail ~"Poisoned exclusive - another task failed inside!";
            }
            rec.failed = true;
            let result = f(&mut rec.data);
            rec.failed = false;
            result
        }
    }
}

// FIXME(#2585) make this a by-move method on the exclusive
#[cfg(stage1)]
fn unwrap_exclusive<T: send>(+arc: Exclusive<T>) -> T {
    let Exclusive { x: x } = arc;
    let inner = unsafe { unwrap_shared_mutable_state(x) };
    let ExData { data: data, _ } = inner;
    data
}

/****************************************************************************
 * Tests
 ****************************************************************************/

#[cfg(test)]
mod tests {

    #[test]
    fn test_reinterpret_cast() {
        assert unsafe { reinterpret_cast(1) } == 1u;
    }

    #[test]
    fn test_bump_box_refcount() {
        unsafe {
            let box = @~"box box box";       // refcount 1
            bump_box_refcount(box);         // refcount 2
            let ptr: *int = transmute(box); // refcount 2
            let _box1: @~str = reinterpret_cast(ptr);
            let _box2: @~str = reinterpret_cast(ptr);
            assert *_box1 == ~"box box box";
            assert *_box2 == ~"box box box";
            // Will destroy _box1 and _box2. Without the bump, this would
            // use-after-free. With too many bumps, it would leak.
        }
    }

    #[test]
    fn test_transmute() {
        unsafe {
            let x = @1;
            let x: *int = transmute(x);
            assert *x == 1;
            let _x: @int = transmute(x);
        }
    }

    #[test]
    fn test_transmute2() {
        unsafe {
            assert transmute(~"L") == ~[76u8, 0u8];
        }
    }

    #[test]
    fn exclusive_arc() {
        let mut futures = ~[];

        let num_tasks = 10u;
        let count = 10u;

        let total = exclusive(~mut 0u);

        for uint::range(0u, num_tasks) |_i| {
            let total = total.clone();
            futures += ~[future::spawn(|| {
                for uint::range(0u, count) |_i| {
                    do total.with |count| {
                        **count += 1u;
                    }
                }
            })];
        };

        for futures.each |f| { f.get() }

        do total.with |total| {
            assert **total == num_tasks * count
        };
    }

    #[test] #[should_fail] #[ignore(cfg(windows))]
    fn exclusive_poison() {
        // Tests that if one task fails inside of an exclusive, subsequent
        // accesses will also fail.
        let x = exclusive(1);
        let x2 = x.clone();
        do task::try {
            do x2.with |one| {
                assert *one == 2;
            }
        };
        do x.with |one| {
            assert *one == 1;
        }
    }

    #[test]
    #[cfg(stage1)]
    fn exclusive_unwrap_basic() {
        let x = exclusive(~~"hello");
        assert unwrap_exclusive(x) == ~~"hello";
    }

    #[test]
    #[cfg(stage1)]
    fn exclusive_unwrap_contended() {
        let x = exclusive(~~"hello");
        let x2 = ~mut some(x.clone());
        do task::spawn {
            let x2 = option::swap_unwrap(x2);
            do x2.with |_hello| { }
            task::yield();
        }
        assert unwrap_exclusive(x) == ~~"hello";

        // Now try the same thing, but with the child task blocking.
        let x = exclusive(~~"hello");
        let x2 = ~mut some(x.clone());
        let mut res = none;
        do task::task().future_result(|+r| res = some(r)).spawn {
            let x2 = option::swap_unwrap(x2);
            assert unwrap_exclusive(x2) == ~~"hello";
        }
        // Have to get rid of our reference before blocking.
        { let _x = move x; } // FIXME(#3161) util::ignore doesn't work here
        let res = option::swap_unwrap(&mut res);
        future::get(&res);
    }

    #[test] #[should_fail] #[ignore(cfg(windows))]
    #[cfg(stage1)]
    fn exclusive_unwrap_conflict() {
        let x = exclusive(~~"hello");
        let x2 = ~mut some(x.clone());
        let mut res = none;
        do task::task().future_result(|+r| res = some(r)).spawn {
            let x2 = option::swap_unwrap(x2);
            assert unwrap_exclusive(x2) == ~~"hello";
        }
        assert unwrap_exclusive(x) == ~~"hello";
        let res = option::swap_unwrap(&mut res);
        future::get(&res);
    }

    #[test] #[ignore(cfg(windows))]
    #[cfg(stage1)]
    fn exclusive_unwrap_deadlock() {
        // This is not guaranteed to get to the deadlock before being killed,
        // but it will show up sometimes, and if the deadlock were not there,
        // the test would nondeterministically fail.
        let result = do task::try {
            // a task that has two references to the same exclusive will
            // deadlock when it unwraps. nothing to be done about that.
            let x = exclusive(~~"hello");
            let x2 = x.clone();
            do task::spawn {
                for 10.times { task::yield(); } // try to let the unwrapper go
                fail; // punt it awake from its deadlock
            }
            let _z = unwrap_exclusive(x);
            do x2.with |_hello| { }
        };
        assert result.is_err();
    }
}
<|endoftext|>


<commit_before>pub mod log_events;
use self::log_events::*;

// http://stackoverflow.com/a/27590832/376138
macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

pub fn parse_line(line: &String) -> Result<LogEvent, &'static str> {
    if line.contains("->") {
        let request = Request::new_from_log_line(line);

        return match request {
            Ok(request) => Ok(LogEvent::Request(request)),
            Err(err) => Err(err),
        };
    }

    if line.contains("<-") {
        let response = Response::new_from_log_line(line);

        return match response {
            Ok(response) => Ok(LogEvent::Response(response)),
            Err(err) => Err(err),
        };
    }

    Err("Line is neither a Request nor a Response")
}

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

    #[test]
    fn test_parse_line_request() {
        let line = "08/Apr/2016:09:58:47 +0200 [02] -> GET /content/some/other.html HTTP/1.1"
            .to_string();

        let event = match parse_line(&line).unwrap() {
            LogEvent::Request(request) => request,
            LogEvent::Response(_) => unreachable!(),
        };

        assert_eq!(event.id, 2);
    }

    #[test]
    fn test_parse_line_response() {
        let line = "08/Apr/2016:09:58:48 +0200 [05] <- 200 text/html 10ms".to_string();

        let event = match parse_line(&line).unwrap() {
            LogEvent::Request(_) => unreachable!(),
            LogEvent::Response(response) => response,
        };

        assert_eq!(event.id, 5);
    }
}
<commit_msg>Module does not do direct stderr output<commit_after>pub mod log_events;
use self::log_events::*;

pub fn parse_line(line: &String) -> Result<LogEvent, &'static str> {
    if line.contains("->") {
        let request = Request::new_from_log_line(line);

        return match request {
            Ok(request) => Ok(LogEvent::Request(request)),
            Err(err) => Err(err),
        };
    }

    if line.contains("<-") {
        let response = Response::new_from_log_line(line);

        return match response {
            Ok(response) => Ok(LogEvent::Response(response)),
            Err(err) => Err(err),
        };
    }

    Err("Line is neither a Request nor a Response")
}

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

    #[test]
    fn test_parse_line_request() {
        let line = "08/Apr/2016:09:58:47 +0200 [02] -> GET /content/some/other.html HTTP/1.1"
            .to_string();

        let event = match parse_line(&line).unwrap() {
            LogEvent::Request(request) => request,
            LogEvent::Response(_) => unreachable!(),
        };

        assert_eq!(event.id, 2);
    }

    #[test]
    fn test_parse_line_response() {
        let line = "08/Apr/2016:09:58:48 +0200 [05] <- 200 text/html 10ms".to_string();

        let event = match parse_line(&line).unwrap() {
            LogEvent::Request(_) => unreachable!(),
            LogEvent::Response(response) => response,
        };

        assert_eq!(event.id, 5);
    }
}
<|endoftext|>


<commit_before>extern crate moho;
extern crate sdl2;
extern crate glm;

use self::moho::input_manager::*;
use self::moho::resource_manager::*;
use self::moho::MohoEngine;

use self::sdl2::keyboard::Keycode;
use self::sdl2::rect;
use self::sdl2::mouse::MouseButton;

use self::glm::ext::normalize_to;

use std::error::Error;

use meteor::Meteor;
use planet::Planet;
use animation::Animation;
use sprite_strip::SpriteStrip;
use explosion::Explosion;

pub struct MasterSmasher<E: MohoEngine> {
    meteor: Meteor<E::Renderer>,
    planets: Vec<Planet<E::Renderer>>,
    background: TextureData<<E::Renderer as Renderer>::Texture>,
    explosion: Option<Explosion<E::Renderer>>,
    input_manager: InputManager<E::EventPump>,
    renderer: ResourceManager<E::Renderer>,
    rects: [rect::Rect; 10],
}

impl<E: MohoEngine> MasterSmasher<E> {
    pub fn new(renderer: ResourceManager<E::Renderer>,
               input_manager: InputManager<E::EventPump>)
               -> Result<Self, Box<Error>> {
        let background_path = "resources/background_game.png";
        let meteor_path = "resources/meteor.png";
        let blue_planet_path = "resources/blue_planet.png";
        let red_planet_path = "resources/red_planet.png";

        let (window_width, window_height) = renderer.output_size()?;
        let background = renderer.load_texture(background_path)?;
        let blue_planet = Planet::new(renderer.load_texture(blue_planet_path)?,
                                      glm::ivec2(400, 300));
        let red_planet = Planet::new(renderer.load_texture(red_planet_path)?,
                                     glm::ivec2(700, 500));
        let meteor = Meteor::new(renderer.load_texture(meteor_path)?,
                                 glm::ivec2(50, 50),
                                 glm::uvec2(window_width, window_height));

        Ok(MasterSmasher {
            meteor: meteor,
            planets: vec![blue_planet, red_planet],
            background: background,
            explosion: None,
            input_manager: input_manager,
            renderer: renderer,
            rects: [rect::Rect::new(0, 0, 5, 5); 10],
        })
    }

    pub fn run(&mut self) -> Result<(), Box<Error>> {
        while self.update() {
            self.draw()?;
        }
        Ok(())
    }

    fn update(&mut self) -> bool {
        if !self.input_manager.update() || self.input_manager.is_key_down(Keycode::Escape) {
            return false;
        }

        if self.meteor.is_launched() {
            if self.input_manager.did_press_key(Keycode::R) {
                self.meteor.restart_at(glm::ivec2(50, 50));
            } else {
                self.update_meteor();
            }
        } else if self.input_manager.did_click_mouse(MouseButton::Left) {
            self.meteor.launch(self.input_manager.mouse_coords());
        } else {
            self.update_launch_vector();
        }

        let explosion_ended = match self.explosion {
            Some(ref mut expl) => !expl.update(),
            None => false,
        };

        if explosion_ended {
            self.explosion = None
        };

        true
    }

    fn draw(&mut self) -> Result<(), Box<Error>> {
        self.renderer.clear();
        self.renderer.draw(&*self.background.texture, None, None, None)?;
        self.meteor.draw(&mut self.renderer)?;
        for planet in &self.planets {
            planet.draw(&mut self.renderer)?;
        }
        if !self.meteor.is_launched() {
            self.renderer.fill_rects(&self.rects)?;
        }
        if let Some(ref expl) = self.explosion {
            expl.draw(&mut self.renderer)?
        }
        self.renderer.present();
        Ok(())
    }

    fn update_meteor(&mut self) {
        self.meteor.update();

        if self.meteor.collides_with(&self.planets) {
            let explosion_path = "resources/explosion_large.png";
            let explosion_texture = self.renderer.load_texture(explosion_path).unwrap();
            let explosion_sprite = SpriteStrip::new(explosion_texture, 8, None);
            let animation = Animation::new(explosion_sprite, 8, false, 80);
            let center = glm::ivec2(self.meteor.center().x as i32, self.meteor.center().y as i32);
            self.explosion = Some(Explosion::new(animation, center));
            self.meteor.restart_at(glm::ivec2(50, 50));
        }
    }

    fn update_launch_vector(&mut self) {
        let mouse_coords = self.input_manager.mouse_coords();
        let mouse_coords = glm::dvec2(mouse_coords.x as f64, mouse_coords.y as f64);
        let distance = mouse_coords - self.meteor.center();
        let offset = self.meteor.radius() + 10.;
        let offset_vector = normalize_to(distance, offset);
        let anchor_point = self.meteor.center() + offset_vector;
        let step = (mouse_coords - anchor_point) / (self.rects.len() as f64);

        for (i, rect) in self.rects.iter_mut().enumerate() {
            let point = anchor_point + (step * i as f64);
            rect.center_on((point.x as i32, point.y as i32));
        }
    }
}
<commit_msg>allow for multiple explosions; implemented my own retain_mut because rust needs one ლ(ಥ Д ಥ )ლ<commit_after>extern crate moho;
extern crate sdl2;
extern crate glm;

use self::moho::input_manager::*;
use self::moho::resource_manager::*;
use self::moho::MohoEngine;

use self::sdl2::keyboard::Keycode;
use self::sdl2::rect;
use self::sdl2::mouse::MouseButton;

use self::glm::ext::normalize_to;

use std::error::Error;

use meteor::Meteor;
use planet::Planet;
use animation::Animation;
use sprite_strip::SpriteStrip;
use explosion::Explosion;

pub struct MasterSmasher<E: MohoEngine> {
    meteor: Meteor<E::Renderer>,
    planets: Vec<Planet<E::Renderer>>,
    background: TextureData<<E::Renderer as Renderer>::Texture>,
    explosions: Vec<Explosion<E::Renderer>>,
    input_manager: InputManager<E::EventPump>,
    renderer: ResourceManager<E::Renderer>,
    rects: [rect::Rect; 10],
}

impl<E: MohoEngine> MasterSmasher<E> {
    pub fn new(renderer: ResourceManager<E::Renderer>,
               input_manager: InputManager<E::EventPump>)
               -> Result<Self, Box<Error>> {
        let background_path = "resources/background_game.png";
        let meteor_path = "resources/meteor.png";
        let blue_planet_path = "resources/blue_planet.png";
        let red_planet_path = "resources/red_planet.png";

        let (window_width, window_height) = renderer.output_size()?;
        let background = renderer.load_texture(background_path)?;
        let blue_planet = Planet::new(renderer.load_texture(blue_planet_path)?,
                                      glm::ivec2(400, 300));
        let red_planet = Planet::new(renderer.load_texture(red_planet_path)?,
                                     glm::ivec2(700, 500));
        let meteor = Meteor::new(renderer.load_texture(meteor_path)?,
                                 glm::ivec2(50, 50),
                                 glm::uvec2(window_width, window_height));

        Ok(MasterSmasher {
            meteor: meteor,
            planets: vec![blue_planet, red_planet],
            background: background,
            explosions: vec![],
            input_manager: input_manager,
            renderer: renderer,
            rects: [rect::Rect::new(0, 0, 5, 5); 10],
        })
    }

    pub fn run(&mut self) -> Result<(), Box<Error>> {
        while self.update() {
            self.draw()?;
        }
        Ok(())
    }

    fn update(&mut self) -> bool {
        if !self.input_manager.update() || self.input_manager.is_key_down(Keycode::Escape) {
            return false;
        }

        if self.meteor.is_launched() {
            if self.input_manager.did_press_key(Keycode::R) {
                self.meteor.restart_at(glm::ivec2(50, 50));
            } else {
                self.update_meteor();
            }
        } else if self.input_manager.did_click_mouse(MouseButton::Left) {
            self.meteor.launch(self.input_manager.mouse_coords());
        } else {
            self.update_launch_vector();
        }

        self.retain_mut(|ref mut e| e.update());

        true
    }

    fn draw(&mut self) -> Result<(), Box<Error>> {
        self.renderer.clear();
        self.renderer.draw(&*self.background.texture, None, None, None)?;
        self.meteor.draw(&mut self.renderer)?;
        for planet in &self.planets {
            planet.draw(&mut self.renderer)?;
        }
        if !self.meteor.is_launched() {
            self.renderer.fill_rects(&self.rects)?;
        }
        for explosion in &self.explosions {
            explosion.draw(&mut self.renderer)?
        }
        self.renderer.present();
        Ok(())
    }

    fn update_meteor(&mut self) {
        self.meteor.update();

        if self.meteor.collides_with(&self.planets) {
            let explosion_path = "resources/explosion_large.png";
            let explosion_texture = self.renderer.load_texture(explosion_path).unwrap();
            let explosion_sprite = SpriteStrip::new(explosion_texture, 8, None);
            let animation = Animation::new(explosion_sprite, 8, false, 80);
            let center = glm::ivec2(self.meteor.center().x as i32, self.meteor.center().y as i32);
            self.explosions.push(Explosion::new(animation, center));
            self.meteor.restart_at(glm::ivec2(50, 50));
        }
    }

    fn update_launch_vector(&mut self) {
        let mouse_coords = self.input_manager.mouse_coords();
        let mouse_coords = glm::dvec2(mouse_coords.x as f64, mouse_coords.y as f64);
        let distance = mouse_coords - self.meteor.center();
        let offset = self.meteor.radius() + 10.;
        let offset_vector = normalize_to(distance, offset);
        let anchor_point = self.meteor.center() + offset_vector;
        let step = (mouse_coords - anchor_point) / (self.rects.len() as f64);

        for (i, rect) in self.rects.iter_mut().enumerate() {
            let point = anchor_point + (step * i as f64);
            rect.center_on((point.x as i32, point.y as i32));
        }
    }

    pub fn retain_mut<F>(&mut self, mut f: F)
        where F: FnMut(&mut Explosion<E::Renderer>) -> bool
    {
        let vector = &mut self.explosions;
        let len = vector.len();
        let mut del = 0;
        {
            let v = &mut **vector;

            for i in 0..len {
                if !f(&mut v[i]) {
                    del += 1;
                } else if del > 0 {
                    v.swap(i - del, i);
                }
            }
        }
        if del > 0 {
            vector.truncate(len - del);
        }
    }
}
<|endoftext|>


<commit_before>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        pwrite_all(&self.file, data, self.heap_id.offset())?;
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(false);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        offset.try_into().unwrap(),
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<commit_msg>Make heap files zero-padded for lexicographic command line ease<commit_after>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        pwrite_all(&self.file, data, self.heap_id.offset())?;
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{:02}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(false);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        offset.try_into().unwrap(),
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<|endoftext|>