<commit_before>//! Game world display utilities.

use std::sync::Arc;
use calx_resource::Resource;
use calx_grid::{Dir12, Dir6};
use world::{Brush, Location, World};
use world::{TerrainQuery, terrain};

/// Surface angle for a visible sprite, used for dynamic lighting.
///
/// ```notrust
///
///      # north #
///     n         n
///    w           e
///
///  # x_         _y #
///   s  -x_   _y-  s
///    w    -*-    e
///       y-   -x
///      # south #
/// ```
#[derive(Copy, Eq, PartialEq, Clone, Debug)]
pub enum Angle {
    Up,
    North,
    XWallBack,
    Northeast,
    East,
    Southeast,
    YWall,
    South,
    XWall,
    Southwest,
    West,
    Northwest,
    YWallBack,
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
/// Draw layer for visual map elements.
pub enum Layer {
    /// Floor sprites, below other map forms.
    Floor,
    /// Blood splatters etc. on floor.
    Decal,
    /// Small items on floor,.
    Items,
    /// Large map objects, walls, most entities etc.
    Object,
    /// Transient effects shown on top of other map content.
    Effect,
    /// Text captions for map elements, on top of everything else.
    Text,
}

pub fn draw_terrain_sprites<F>(w: &World, loc: Location, mut draw: F)
    where F: FnMut(Layer, Angle, &Resource<Brush>, usize)
{
    use self::Angle::*;

    let terrain = w.terrain(loc);
    let kernel = Kernel::new(w, loc);

    match terrain.form {
        terrain::Form::Void | terrain::Form::Floor => {
            draw(Layer::Floor, Up, &terrain.brush, 0);
        }
        terrain::Form::Gate => {
            if let Some(d12) = Dir12::away_from(&kernel.void_mask()) {
                draw(Layer::Floor, Up, &terrain.brush, d12 as usize + 1);
            } else {
                draw(Layer::Floor, Up, &terrain.brush, 0);
            }
        }
        terrain::Form::Prop => {
            draw(Layer::Object, South, &terrain.brush, 0);
        }
        terrain::Form::Blob => {
            // This part gets a little tricky. Basic idea is that there's an inner pointy-top
            // hex core and the blob hull will snap to that instead of the outer flat-top hex
            // edge if neither adjacent face to the outer hex vertex is connected to another
            // blob.
            //
            // Based on how the sprites split up, the processing is done in four vertical
            // segments.

            let faces = kernel.blob_faces();

            // Do we snap to the outer vertices?
            let ne_vertex = !faces[0] || !faces[1];
            let e_vertex = !faces[1] || !faces[2];
            let se_vertex = !faces[2] || !faces[3];
            let sw_vertex = !faces[3] || !faces[4];
            let w_vertex = !faces[4] || !faces[5];
            let nw_vertex = !faces[5] || !faces[0];

            // Segment 2, middle left
            {
                if faces[0] {
                    if nw_vertex && ne_vertex {
                        draw(Layer::Object, North, &terrain.brush, 7);
                    } else if nw_vertex {
                        draw(Layer::Object, XWallBack, &terrain.brush, 15);
                    } else {
                        draw(Layer::Object, YWallBack, &terrain.brush, 11);
                    }
                }
                if faces[3] {
                    if sw_vertex && se_vertex {
                        draw(Layer::Object, South, &terrain.brush, 19);
                    } else if sw_vertex {
                        draw(Layer::Object, YWall, &terrain.brush, 23);
                    } else {
                        draw(Layer::Object, XWall, &terrain.brush, 27);
                    }
                }
            }

            // Segment 3, middle right
            {
                if faces[0] {
                    if ne_vertex && nw_vertex {
                        draw(Layer::Object, North, &terrain.brush, 8);
                    } else if ne_vertex {
                        draw(Layer::Object, YWallBack, &terrain.brush, 12);
                    } else {
                        draw(Layer::Object, XWallBack, &terrain.brush, 16);
                    }
                }
                if faces[3] {
                    if se_vertex && sw_vertex {
                        draw(Layer::Object, South, &terrain.brush, 20);
                    } else if se_vertex {
                        draw(Layer::Object, XWall, &terrain.brush, 28);
                    } else {
                        draw(Layer::Object, YWall, &terrain.brush, 24);
                    }
                }
            }

            // The side segments need to come after the middle
            // segments so that the vertical edges can overwrite the
            // middle segment pixels.

            // Segment 1, left edge
            {
                if w_vertex {
                    if faces[5] {
                        if nw_vertex {
                            draw(Layer::Object, Northwest, &terrain.brush, 6);
                        } else {
                            draw(Layer::Object, YWallBack, &terrain.brush, 10);
                        }
                    }

                    if faces[4] {
                        if sw_vertex {
                            draw(Layer::Object, Southwest, &terrain.brush, 18);
                        } else {
                            draw(Layer::Object, XWall, &terrain.brush, 26);
                        }
                    }
                } else {
                    // Draw the left vertical line.
                    draw(Layer::Object, West, &terrain.brush, 2);
                    if !faces[0] {
                        draw(Layer::Object, West, &terrain.brush, 0);
                    }
                    if !faces[3] {
                        draw(Layer::Object, West, &terrain.brush, 4);
                    }
                }
            }

            // Segment 4, right edge
            {
                if e_vertex {
                    if faces[1] {
                        if ne_vertex {
                            draw(Layer::Object, Northeast, &terrain.brush, 9);
                        } else {
                            draw(Layer::Object, XWallBack, &terrain.brush, 17);
                        }
                    }

                    if faces[2] {
                        if se_vertex {
                            draw(Layer::Object, Southeast, &terrain.brush, 21);
                        } else {
                            draw(Layer::Object, YWall, &terrain.brush, 25);
                        }
                    }
                } else {
                    // Draw the right vertical line.
                    draw(Layer::Object, East, &terrain.brush, 3);
                    if !faces[0] {
                        draw(Layer::Object, East, &terrain.brush, 1);
                    }
                    if !faces[3] {
                        draw(Layer::Object, East, &terrain.brush, 5);
                    }
                }
            }
        }
        terrain::Form::Wall => {
            let extends = kernel.wall_extends();
            if extends[0] {
                draw(Layer::Object, XWall, &terrain.brush, 2);
            } else {
                draw(Layer::Object, XWall, &terrain.brush, 0);
            }
            if extends[1] {
                draw(Layer::Object, YWall, &terrain.brush, 3);
            } else {
                draw(Layer::Object, YWall, &terrain.brush, 1);
            }
        }
    }
    // TODO: Generate sprites for entities, with tweening state projection.
    // TODO: Generate special effect sprites grounded on this location.
}

#[derive(Clone)]
pub struct Kernel {
    pub n: Arc<terrain::Tile>,
    pub ne: Arc<terrain::Tile>,
    pub nw: Arc<terrain::Tile>,
    pub center: Arc<terrain::Tile>,
    pub se: Arc<terrain::Tile>,
    pub sw: Arc<terrain::Tile>,
    pub s: Arc<terrain::Tile>,
}

fn neighbor(w: &World, loc: Location, dir: Dir6) -> Arc<terrain::Tile> {
    let loc = w.visible_portal(loc + dir.to_v2()).unwrap_or(loc + dir.to_v2());
    w.terrain(loc)
}

impl Kernel {
    pub fn new(w: &World, loc: Location) -> Kernel {
        Kernel {
            n: neighbor(w, loc, Dir6::North),
            ne: neighbor(w, loc, Dir6::Northeast),
            nw: neighbor(w, loc, Dir6::Northwest),
            center: w.terrain(loc),
            se: neighbor(w, loc, Dir6::Southeast),
            sw: neighbor(w, loc, Dir6::Southwest),
            s: neighbor(w, loc, Dir6::South),
        }
    }

    /// Bool is true if left/right half of wall should be extended.
    pub fn wall_extends(&self) -> [bool; 2] { [self.nw.is_wall(), self.ne.is_wall()] }

    /// Bool is true if n/ne/se/s/sw/nw face of block is facing open air.
    pub fn blob_faces(&self) -> [bool; 6] {
        // Because they work a bit differently visually, back-side faces
        // are not drawn if there is any hull touching, front is only
        // not drawn if there's another blob.
        [!self.n.is_hull(),
         !self.ne.is_hull(),
         !self.se.is_blob(),
         !self.s.is_blob(),
         !self.sw.is_blob(),
         !self.nw.is_hull()]
    }

    /// Mark neighbors that are not void terrain as true.
    pub fn void_mask(&self) -> [bool; 6] {
        [self.n.form != terrain::Form::Void,
         self.ne.form != terrain::Form::Void,
         self.se.form != terrain::Form::Void,
         self.s.form != terrain::Form::Void,
         self.sw.form != terrain::Form::Void,
         self.nw.form != terrain::Form::Void]
    }
}
<commit_msg>Nicer looking blob hulling<commit_after>//! Game world display utilities.

use std::sync::Arc;
use calx_resource::Resource;
use calx_grid::{Dir12, Dir6};
use world::{Brush, Location, World};
use world::{TerrainQuery, terrain};

/// Surface angle for a visible sprite, used for dynamic lighting.
///
/// ```notrust
///
///      # north #
///     n         n
///    w           e
///
///  # x_         _y #
///   s  -x_   _y-  s
///    w    -*-    e
///       y-   -x
///      # south #
/// ```
#[derive(Copy, Eq, PartialEq, Clone, Debug)]
pub enum Angle {
    Up,
    North,
    XWallBack,
    Northeast,
    East,
    Southeast,
    YWall,
    South,
    XWall,
    Southwest,
    West,
    Northwest,
    YWallBack,
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
/// Draw layer for visual map elements.
pub enum Layer {
    /// Floor sprites, below other map forms.
    Floor,
    /// Blood splatters etc. on floor.
    Decal,
    /// Small items on floor,.
    Items,
    /// Large map objects, walls, most entities etc.
    Object,
    /// Transient effects shown on top of other map content.
    Effect,
    /// Text captions for map elements, on top of everything else.
    Text,
}

pub fn draw_terrain_sprites<F>(w: &World, loc: Location, mut draw: F)
    where F: FnMut(Layer, Angle, &Resource<Brush>, usize)
{
    use self::Angle::*;

    let terrain = w.terrain(loc);
    let kernel = Kernel::new(w, loc);

    match terrain.form {
        terrain::Form::Void | terrain::Form::Floor => {
            draw(Layer::Floor, Up, &terrain.brush, 0);
        }
        terrain::Form::Gate => {
            if let Some(d12) = Dir12::away_from(&kernel.void_mask()) {
                draw(Layer::Floor, Up, &terrain.brush, d12 as usize + 1);
            } else {
                draw(Layer::Floor, Up, &terrain.brush, 0);
            }
        }
        terrain::Form::Prop => {
            draw(Layer::Object, South, &terrain.brush, 0);
        }
        terrain::Form::Blob => {
            // This part gets a little tricky. Basic idea is that there's an inner pointy-top
            // hex core and the blob hull will snap to that instead of the outer flat-top hex
            // edge if neither adjacent face to the outer hex vertex is connected to another
            // blob.
            //
            // Based on how the sprites split up, the processing is done in four vertical
            // segments.

            let faces = kernel.blob_faces();

            // Do we snap to the outer vertices?
            let ne_vertex = !faces[0] || !faces[1];
            let e_vertex = !faces[1] || !faces[2];
            let se_vertex = !faces[2] || !faces[3];
            let sw_vertex = !faces[3] || !faces[4];
            let w_vertex = !faces[4] || !faces[5];
            let nw_vertex = !faces[5] || !faces[0];

            // Segment 2, middle left
            {
                if faces[0] {
                    if nw_vertex && ne_vertex {
                        draw(Layer::Object, North, &terrain.brush, 7);
                    } else if nw_vertex {
                        draw(Layer::Object, XWallBack, &terrain.brush, 15);
                    } else {
                        draw(Layer::Object, YWallBack, &terrain.brush, 11);
                    }
                }
                if faces[3] {
                    if sw_vertex && se_vertex {
                        draw(Layer::Object, South, &terrain.brush, 19);
                    } else if sw_vertex {
                        draw(Layer::Object, YWall, &terrain.brush, 23);
                    } else {
                        draw(Layer::Object, XWall, &terrain.brush, 27);
                    }
                }
            }

            // Segment 3, middle right
            {
                if faces[0] {
                    if ne_vertex && nw_vertex {
                        draw(Layer::Object, North, &terrain.brush, 8);
                    } else if ne_vertex {
                        draw(Layer::Object, YWallBack, &terrain.brush, 12);
                    } else {
                        draw(Layer::Object, XWallBack, &terrain.brush, 16);
                    }
                }
                if faces[3] {
                    if se_vertex && sw_vertex {
                        draw(Layer::Object, South, &terrain.brush, 20);
                    } else if se_vertex {
                        draw(Layer::Object, XWall, &terrain.brush, 28);
                    } else {
                        draw(Layer::Object, YWall, &terrain.brush, 24);
                    }
                }
            }

            // The side segments need to come after the middle
            // segments so that the vertical edges can overwrite the
            // middle segment pixels.

            // Segment 1, left edge
            {
                if w_vertex {
                    if faces[5] {
                        if nw_vertex {
                            draw(Layer::Object, Northwest, &terrain.brush, 6);
                        } else {
                            draw(Layer::Object, YWallBack, &terrain.brush, 10);
                        }
                    }

                    if faces[4] {
                        if sw_vertex {
                            draw(Layer::Object, Southwest, &terrain.brush, 18);
                        } else {
                            draw(Layer::Object, XWall, &terrain.brush, 26);
                        }
                    }
                } else {
                    // Draw the left vertical line.
                    draw(Layer::Object, West, &terrain.brush, 2);
                    if !faces[0] {
                        draw(Layer::Object, West, &terrain.brush, 0);
                    }
                    if !faces[3] {
                        draw(Layer::Object, West, &terrain.brush, 4);
                    }
                }
            }

            // Segment 4, right edge
            {
                if e_vertex {
                    if faces[1] {
                        if ne_vertex {
                            draw(Layer::Object, Northeast, &terrain.brush, 9);
                        } else {
                            draw(Layer::Object, XWallBack, &terrain.brush, 17);
                        }
                    }

                    if faces[2] {
                        if se_vertex {
                            draw(Layer::Object, Southeast, &terrain.brush, 21);
                        } else {
                            draw(Layer::Object, YWall, &terrain.brush, 25);
                        }
                    }
                } else {
                    // Draw the right vertical line.
                    draw(Layer::Object, East, &terrain.brush, 3);
                    if !faces[0] {
                        draw(Layer::Object, East, &terrain.brush, 1);
                    }
                    if !faces[3] {
                        draw(Layer::Object, East, &terrain.brush, 5);
                    }
                }
            }
        }
        terrain::Form::Wall => {
            let extends = kernel.wall_extends();
            if extends[0] {
                draw(Layer::Object, XWall, &terrain.brush, 2);
            } else {
                draw(Layer::Object, XWall, &terrain.brush, 0);
            }
            if extends[1] {
                draw(Layer::Object, YWall, &terrain.brush, 3);
            } else {
                draw(Layer::Object, YWall, &terrain.brush, 1);
            }
        }
    }
    // TODO: Generate sprites for entities, with tweening state projection.
    // TODO: Generate special effect sprites grounded on this location.
}

#[derive(Clone)]
pub struct Kernel {
    pub n: Arc<terrain::Tile>,
    pub ne: Arc<terrain::Tile>,
    pub nw: Arc<terrain::Tile>,
    pub center: Arc<terrain::Tile>,
    pub se: Arc<terrain::Tile>,
    pub sw: Arc<terrain::Tile>,
    pub s: Arc<terrain::Tile>,
}

fn neighbor(w: &World, loc: Location, dir: Dir6) -> Arc<terrain::Tile> {
    let loc = w.visible_portal(loc + dir.to_v2()).unwrap_or(loc + dir.to_v2());
    w.terrain(loc)
}

impl Kernel {
    pub fn new(w: &World, loc: Location) -> Kernel {
        Kernel {
            n: neighbor(w, loc, Dir6::North),
            ne: neighbor(w, loc, Dir6::Northeast),
            nw: neighbor(w, loc, Dir6::Northwest),
            center: w.terrain(loc),
            se: neighbor(w, loc, Dir6::Southeast),
            sw: neighbor(w, loc, Dir6::Southwest),
            s: neighbor(w, loc, Dir6::South),
        }
    }

    /// Bool is true if left/right half of wall should be extended.
    pub fn wall_extends(&self) -> [bool; 2] { [self.nw.is_wall(), self.ne.is_wall()] }

    /// Bool is true if n/ne/se/s/sw/nw face of block is facing open air.
    pub fn blob_faces(&self) -> [bool; 6] {
        [!self.n.is_hull(),
         !self.ne.is_hull(),
         !self.se.is_hull(),
         !self.s.is_hull(),
         !self.sw.is_hull(),
         !self.nw.is_hull()]
    }

    /// Mark neighbors that are not void terrain as true.
    pub fn void_mask(&self) -> [bool; 6] {
        [self.n.form != terrain::Form::Void,
         self.ne.form != terrain::Form::Void,
         self.se.form != terrain::Form::Void,
         self.s.form != terrain::Form::Void,
         self.sw.form != terrain::Form::Void,
         self.nw.form != terrain::Form::Void]
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use std::mem::transmute;
use std::slice;

pub struct Writer {
    enc: [u8; 100]
}

extern {
    fn od_ec_enc_init(enc: *mut libc::c_void, size: u32);
    fn od_ec_enc_done(enc: *mut libc::c_void, nbytes: *mut u32) -> *const u8;
    fn od_ec_encode_cdf_q15(enc: *mut libc::c_void, s: libc::c_int, cdf: *const u16,
                            nsyms: libc::c_int);
    fn od_ec_encode_bool_q15(enc: *mut libc::c_void, val: libc::c_int, f: libc::c_uint);
}

impl Writer {
    pub fn new() -> Writer {
        let enc = [0;100];
        unsafe {
            od_ec_enc_init(transmute(&enc), 1024);
        }
        Writer { enc: enc }
    }
    pub fn done(&mut self) -> &[u8] {
        let mut nbytes: u32 = 0;
        unsafe {
            let b = od_ec_enc_done(transmute(&self.enc), &mut nbytes);
            slice::from_raw_parts(b, nbytes as usize)
        }
    }
    pub fn cdf(&mut self, s: u32, cdf: &[u16], nsyms: u32) {
        unsafe {
            od_ec_encode_cdf_q15(transmute(&self.enc), s as libc::c_int, cdf.as_ptr(), nsyms as libc::c_int);
        }
    }
    pub fn bool(&mut self, val: bool, f: u16) {
        unsafe {
            od_ec_encode_bool_q15(transmute(&self.enc), val as libc::c_int, f as libc::c_uint);
        }
    }
    fn update_cdf(cdf: &mut [u16], val: u32, nsymbs: usize) {
        let rate = 4 + if cdf[nsymbs] > 31 { 1 } else { 0 } + (31 ^ (nsymbs as u32).leading_zeros());
        let rate2 = 5;
        let mut tmp: i32;
        let diff: i32;
        let tmp0 = 1 << rate2;
        tmp = 32768 - tmp0;
        diff = ((32768 - ((nsymbs as i32) << rate2)) >> rate) << rate;
        for i in 0..(nsymbs - 1) {
            if i as u32 == val {
                tmp -= diff;
            }
            cdf[i as usize] = ((cdf[i as usize] as i32) + ((tmp - (cdf[i as usize] as i32)) >> rate)) as u16;
            tmp -= tmp0;
        }
        if cdf[nsymbs] < 32 {
            cdf[nsymbs] += 1;
        }
    }
    pub fn symbol(&mut self, s: u32, cdf: &mut [u16], nsymbs: usize) {
        self.cdf(s, cdf, nsymbs as u32);
        Writer::update_cdf(cdf, s, nsymbs);
    }
}
<commit_msg>Don't bother initializing ec struct.<commit_after>extern crate libc;

use std::mem::transmute;
use std::mem::uninitialized;
use std::slice;

pub struct Writer {
    enc: [u8; 100]
}

extern {
    fn od_ec_enc_init(enc: *mut libc::c_void, size: u32);
    fn od_ec_enc_done(enc: *mut libc::c_void, nbytes: *mut u32) -> *const u8;
    fn od_ec_encode_cdf_q15(enc: *mut libc::c_void, s: libc::c_int, cdf: *const u16,
                            nsyms: libc::c_int);
    fn od_ec_encode_bool_q15(enc: *mut libc::c_void, val: libc::c_int, f: libc::c_uint);
}

impl Writer {
    pub fn new() -> Writer {
        unsafe {
            let enc = uninitialized();
            od_ec_enc_init(transmute(&enc), 1024);
            Writer { enc: enc }
        }
    }
    pub fn done(&mut self) -> &[u8] {
        let mut nbytes: u32 = 0;
        unsafe {
            let b = od_ec_enc_done(transmute(&self.enc), &mut nbytes);
            slice::from_raw_parts(b, nbytes as usize)
        }
    }
    pub fn cdf(&mut self, s: u32, cdf: &[u16], nsyms: u32) {
        unsafe {
            od_ec_encode_cdf_q15(transmute(&self.enc), s as libc::c_int, cdf.as_ptr(), nsyms as libc::c_int);
        }
    }
    pub fn bool(&mut self, val: bool, f: u16) {
        unsafe {
            od_ec_encode_bool_q15(transmute(&self.enc), val as libc::c_int, f as libc::c_uint);
        }
    }
    fn update_cdf(cdf: &mut [u16], val: u32, nsymbs: usize) {
        let rate = 4 + if cdf[nsymbs] > 31 { 1 } else { 0 } + (31 ^ (nsymbs as u32).leading_zeros());
        let rate2 = 5;
        let mut tmp: i32;
        let diff: i32;
        let tmp0 = 1 << rate2;
        tmp = 32768 - tmp0;
        diff = ((32768 - ((nsymbs as i32) << rate2)) >> rate) << rate;
        for i in 0..(nsymbs - 1) {
            if i as u32 == val {
                tmp -= diff;
            }
            cdf[i as usize] = ((cdf[i as usize] as i32) + ((tmp - (cdf[i as usize] as i32)) >> rate)) as u16;
            tmp -= tmp0;
        }
        if cdf[nsymbs] < 32 {
            cdf[nsymbs] += 1;
        }
    }
    pub fn symbol(&mut self, s: u32, cdf: &mut [u16], nsymbs: usize) {
        self.cdf(s, cdf, nsymbs as u32);
        Writer::update_cdf(cdf, s, nsymbs);
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, stdin};
use std::fmt;
use std::fs::File;
use std::sync::Mutex;

use vm::{self, Result, Variants};
use vm::gc::{Gc, Traverseable};
use vm::types::*;
use vm::thread::ThreadInternal;
use vm::thread::Thread;
use vm::api::{Array, FunctionRef, Generic, Hole, VmType, Getable, OpaqueValue, IO, WithVM, Userdata};
use vm::api::generic::{A, B};
use vm::stack::StackFrame;
use vm::internal::ValuePrinter;

use vm::internal::Value;

use super::{Compiler, Error};

fn print(s: &str) -> IO<()> {
    print!("{}", s);
    IO::Value(())
}

fn println(s: &str) -> IO<()> {
    println!("{}", s);
    IO::Value(())
}

struct GluonFile(Mutex<File>);

impl Userdata for GluonFile {}

impl fmt::Debug for GluonFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "File")
    }
}

impl VmType for GluonFile {
    type Type = GluonFile;
}

impl Traverseable for GluonFile {
    fn traverse(&self, _: &mut Gc) {}
}

fn open_file(s: &str) -> IO<GluonFile> {
    match File::open(s) {
        Ok(f) => IO::Value(GluonFile(Mutex::new(f))),
        Err(err) => IO::Exception(format!("{}", err)),
    }
}

fn read_file<'vm>(file: WithVM<'vm, &GluonFile>, count: usize) -> IO<Array<'vm, u8>> {
    let WithVM { vm, value: file } = file;
    let mut file = file.0.lock().unwrap();
    let mut buffer = Vec::with_capacity(count);
    unsafe {
        buffer.set_len(count);
        match file.read(&mut *buffer) {
            Ok(bytes_read) => {
                let value = {
                    let mut context = vm.context();
                    match context.alloc(&buffer[..bytes_read]) {
                        Ok(value) => value,
                        Err(err) => return IO::Exception(format!("{}", err)),
                    }
                };
                IO::Value(Getable::from_value(vm, Variants::new(&Value::Array(value)))
                    .expect("Array"))
            }
            Err(err) => IO::Exception(format!("{}", err)),
        }
    }
}

fn read_file_to_string(s: &str) -> IO<String> {
    let mut buffer = String::new();
    match File::open(s).and_then(|mut file| file.read_to_string(&mut buffer)) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

fn read_char() -> IO<char> {
    match stdin().bytes().next() {
        Some(result) => {
            match result {
                Ok(b) => {
                    ::std::char::from_u32(b as u32)
                        .map(IO::Value)
                        .unwrap_or_else(|| IO::Exception("Not a valid char".into()))
                }
                Err(err) => IO::Exception(format!("{}", err)),
            }
        }
        None => IO::Exception("No read".into()),
    }
}

fn read_line() -> IO<String> {
    let mut buffer = String::new();
    match stdin().read_line(&mut buffer) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

/// IO a -> (String -> IO a) -> IO a
fn catch<'vm>(action: OpaqueValue<&'vm Thread, IO<A>>,
              mut catch: FunctionRef<fn(String) -> IO<Generic<A>>>)
              -> IO<Generic<A>> {
    let vm = action.vm();
    let frame_level = vm.context().stack.get_frames().len();
    let mut action: FunctionRef<fn(()) -> Generic<A>> =
        unsafe { Getable::from_value(vm, Variants::new(&action.get_value())).unwrap() };
    let result = action.call(());
    match result {
        Ok(value) => IO::Value(value),
        Err(err) => {
            {
                let mut context = vm.context();
                let mut stack = StackFrame::current(&mut context.stack);
                while stack.stack.get_frames().len() > frame_level {
                    if stack.exit_scope().is_err() {
                        return IO::Exception("Unknown error".into());
                    }
                }
            }
            match catch.call(format!("{}", err)) {
                Ok(value) => value,
                Err(err) => IO::Exception(format!("{}", err)),
            }
        }
    }
}

fn clear_frames(err: Error, frame_level: usize, mut stack: StackFrame) -> IO<String> {
    let fmt = match err {
        Error::VM(vm::Error::Panic(_)) => {
            let trace = stack.stack.stacktrace(frame_level);
            format!("{}\n{}", err, trace)
        }
        _ => format!("{}", err),
    };
    while stack.stack.get_frames().len() > frame_level {
        if stack.exit_scope().is_err() {
            return IO::Exception(fmt);
        }
    }
    IO::Exception(fmt)
}

fn run_expr(WithVM { vm, value: expr }: WithVM<&str>) -> IO<String> {
    let frame_level = vm.context().stack.get_frames().len();
    let run_result = Compiler::new().run_expr::<OpaqueValue<&Thread, Hole>>(vm, "<top>", expr);
    let mut context = vm.context();
    let stack = StackFrame::current(&mut context.stack);
    match run_result {
        Ok((value, typ)) => {
            let env = vm.global_env().get_env();
            unsafe {
                IO::Value(format!("{} : {}",
                                  ValuePrinter::new(&*env, &typ, value.get_value()).width(80),
                                  typ))
            }
        }
        Err(err) => clear_frames(err, frame_level, stack),
    }
}

fn load_script(WithVM { vm, value: name }: WithVM<&str>, expr: &str) -> IO<String> {
    let frame_level = vm.context().stack.get_frames().len();
    let run_result = Compiler::new().load_script_async(vm, name, expr).sync_or_error();
    let mut context = vm.context();
    let stack = StackFrame::current(&mut context.stack);
    match run_result {
        Ok(()) => IO::Value(format!("Loaded {}", name)),
        Err(err) => clear_frames(err, frame_level, stack),
    }
}

pub fn load(vm: &Thread) -> Result<()> {

    vm.register_type::<GluonFile>("File", &[])?;

    // io_flat_map f m : (a -> IO b) -> IO a -> IO b
    //     = f (m ())
    let io_flat_map = vec![// [f, m, ()]       Initial stack
                           Call(1), // [f, m_ret]       Call m ()
                           PushInt(0), // [f, m_ret, ()]   Add a dummy argument ()
                           TailCall(2) /* [f_ret]          Call f m_ret () */];
    let io_flat_map_type = <fn (fn (A) -> IO<B>, IO<A>) -> IO<B> as VmType>::make_type(vm);
    vm.add_bytecode("io_flat_map", io_flat_map_type, 3, io_flat_map)?;

    vm.add_bytecode("io_pure",
                      <fn(A) -> IO<A> as VmType>::make_type(vm),
                      2,
                      vec![Pop(1)])?;
    // IO functions
    vm.define_global("io",
                       record!(
        open_file => primitive!(1 open_file),
        read_file => primitive!(2 read_file),
        read_file_to_string => primitive!(1 read_file_to_string),
        read_char => primitive!(0 read_char),
        read_line => primitive!(0 read_line),
        print => primitive!(1 print),
        println => primitive!(1 println),
        catch =>
            primitive!(2 catch),
        run_expr => primitive!(1 run_expr),
        load_script => primitive!(2 load_script)
    ))?;
    Ok(())
}
<commit_msg>fix: Run io actions in io.run_expr<commit_after>use std::io::{Read, stdin};
use std::fmt;
use std::fs::File;
use std::sync::Mutex;

use vm::{self, Result, Variants};
use vm::gc::{Gc, Traverseable};
use vm::types::*;
use vm::thread::ThreadInternal;
use vm::thread::Thread;
use vm::api::{Array, FunctionRef, Generic, Hole, VmType, Getable, OpaqueValue, IO, WithVM, Userdata};
use vm::api::generic::{A, B};
use vm::stack::StackFrame;
use vm::internal::ValuePrinter;

use vm::internal::Value;

use super::{Compiler, Error};

fn print(s: &str) -> IO<()> {
    print!("{}", s);
    IO::Value(())
}

fn println(s: &str) -> IO<()> {
    println!("{}", s);
    IO::Value(())
}

struct GluonFile(Mutex<File>);

impl Userdata for GluonFile {}

impl fmt::Debug for GluonFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "File")
    }
}

impl VmType for GluonFile {
    type Type = GluonFile;
}

impl Traverseable for GluonFile {
    fn traverse(&self, _: &mut Gc) {}
}

fn open_file(s: &str) -> IO<GluonFile> {
    match File::open(s) {
        Ok(f) => IO::Value(GluonFile(Mutex::new(f))),
        Err(err) => IO::Exception(format!("{}", err)),
    }
}

fn read_file<'vm>(file: WithVM<'vm, &GluonFile>, count: usize) -> IO<Array<'vm, u8>> {
    let WithVM { vm, value: file } = file;
    let mut file = file.0.lock().unwrap();
    let mut buffer = Vec::with_capacity(count);
    unsafe {
        buffer.set_len(count);
        match file.read(&mut *buffer) {
            Ok(bytes_read) => {
                let value = {
                    let mut context = vm.context();
                    match context.alloc(&buffer[..bytes_read]) {
                        Ok(value) => value,
                        Err(err) => return IO::Exception(format!("{}", err)),
                    }
                };
                IO::Value(Getable::from_value(vm, Variants::new(&Value::Array(value)))
                    .expect("Array"))
            }
            Err(err) => IO::Exception(format!("{}", err)),
        }
    }
}

fn read_file_to_string(s: &str) -> IO<String> {
    let mut buffer = String::new();
    match File::open(s).and_then(|mut file| file.read_to_string(&mut buffer)) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

fn read_char() -> IO<char> {
    match stdin().bytes().next() {
        Some(result) => {
            match result {
                Ok(b) => {
                    ::std::char::from_u32(b as u32)
                        .map(IO::Value)
                        .unwrap_or_else(|| IO::Exception("Not a valid char".into()))
                }
                Err(err) => IO::Exception(format!("{}", err)),
            }
        }
        None => IO::Exception("No read".into()),
    }
}

fn read_line() -> IO<String> {
    let mut buffer = String::new();
    match stdin().read_line(&mut buffer) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

/// IO a -> (String -> IO a) -> IO a
fn catch<'vm>(action: OpaqueValue<&'vm Thread, IO<A>>,
              mut catch: FunctionRef<fn(String) -> IO<Generic<A>>>)
              -> IO<Generic<A>> {
    let vm = action.vm();
    let frame_level = vm.context().stack.get_frames().len();
    let mut action: FunctionRef<fn(()) -> Generic<A>> =
        unsafe { Getable::from_value(vm, Variants::new(&action.get_value())).unwrap() };
    let result = action.call(());
    match result {
        Ok(value) => IO::Value(value),
        Err(err) => {
            {
                let mut context = vm.context();
                let mut stack = StackFrame::current(&mut context.stack);
                while stack.stack.get_frames().len() > frame_level {
                    if stack.exit_scope().is_err() {
                        return IO::Exception("Unknown error".into());
                    }
                }
            }
            match catch.call(format!("{}", err)) {
                Ok(value) => value,
                Err(err) => IO::Exception(format!("{}", err)),
            }
        }
    }
}

fn clear_frames(err: Error, frame_level: usize, mut stack: StackFrame) -> IO<String> {
    let fmt = match err {
        Error::VM(vm::Error::Panic(_)) => {
            let trace = stack.stack.stacktrace(frame_level);
            format!("{}\n{}", err, trace)
        }
        _ => format!("{}", err),
    };
    while stack.stack.get_frames().len() > frame_level {
        if stack.exit_scope().is_err() {
            return IO::Exception(fmt);
        }
    }
    IO::Exception(fmt)
}

fn run_expr(WithVM { vm, value: expr }: WithVM<&str>) -> IO<String> {
    let frame_level = vm.context().stack.get_frames().len();
    let run_result = Compiler::new().run_io_expr::<OpaqueValue<&Thread, Hole>>(vm, "<top>", expr);
    let mut context = vm.context();
    let stack = StackFrame::current(&mut context.stack);
    match run_result {
        Ok((value, typ)) => {
            let env = vm.global_env().get_env();
            unsafe {
                IO::Value(format!("{} : {}",
                                  ValuePrinter::new(&*env, &typ, value.get_value()).width(80),
                                  typ))
            }
        }
        Err(err) => clear_frames(err, frame_level, stack),
    }
}

fn load_script(WithVM { vm, value: name }: WithVM<&str>, expr: &str) -> IO<String> {
    let frame_level = vm.context().stack.get_frames().len();
    let run_result = Compiler::new().load_script_async(vm, name, expr).sync_or_error();
    let mut context = vm.context();
    let stack = StackFrame::current(&mut context.stack);
    match run_result {
        Ok(()) => IO::Value(format!("Loaded {}", name)),
        Err(err) => clear_frames(err, frame_level, stack),
    }
}

pub fn load(vm: &Thread) -> Result<()> {

    vm.register_type::<GluonFile>("File", &[])?;

    // io_flat_map f m : (a -> IO b) -> IO a -> IO b
    //     = f (m ())
    let io_flat_map = vec![// [f, m, ()]       Initial stack
                           Call(1), // [f, m_ret]       Call m ()
                           PushInt(0), // [f, m_ret, ()]   Add a dummy argument ()
                           TailCall(2) /* [f_ret]          Call f m_ret () */];
    let io_flat_map_type = <fn (fn (A) -> IO<B>, IO<A>) -> IO<B> as VmType>::make_type(vm);
    vm.add_bytecode("io_flat_map", io_flat_map_type, 3, io_flat_map)?;

    vm.add_bytecode("io_pure",
                      <fn(A) -> IO<A> as VmType>::make_type(vm),
                      2,
                      vec![Pop(1)])?;
    // IO functions
    vm.define_global("io",
                       record!(
        open_file => primitive!(1 open_file),
        read_file => primitive!(2 read_file),
        read_file_to_string => primitive!(1 read_file_to_string),
        read_char => primitive!(0 read_char),
        read_line => primitive!(0 read_line),
        print => primitive!(1 print),
        println => primitive!(1 println),
        catch =>
            primitive!(2 catch),
        run_expr => primitive!(1 run_expr),
        load_script => primitive!(2 load_script)
    ))?;
    Ok(())
}
<|endoftext|>


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

//! Interfaces to the operating system provided random number
//! generators.

use std::io;
use Rng;

/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`, or from `getrandom(2)` system call if available.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
/// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
/// - PNaCl: calls into the `nacl-irt-random-0.1` IRT interface.
///
/// This does not block.
pub struct OsRng(imp::OsRng);

impl OsRng {
    /// Create a new `OsRng`.
    pub fn new() -> io::Result<OsRng> {
        imp::OsRng::new().map(OsRng)
    }
}

impl Rng for OsRng {
    fn next_u32(&mut self) -> u32 { self.0.next_u32() }
    fn next_u64(&mut self) -> u64 { self.0.next_u64() }
    fn fill_bytes(&mut self, v: &mut [u8]) { self.0.fill_bytes(v) }
}

#[cfg(all(unix, not(target_os = "ios"),
          not(target_os = "nacl")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use std::io;
    use std::fs::File;
    use Rng;
    use read::ReadRng;
    use std::mem;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(target_arch = "aarch64")]
        const NR_GETRANDOM: libc::c_long = 278;
        #[cfg(target_arch = "powerpc")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0u8; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0u8; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn is_getrandom_available() -> bool {
        use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
        use std::sync::{Once, ONCE_INIT};

        static CHECKER: Once = ONCE_INIT;
        static AVAILABLE: AtomicBool = ATOMIC_BOOL_INIT;

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn is_getrandom_available() -> bool { false }

    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReadRng(ReadRng<File>),
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReadRng::new(reader);

            Ok(OsRng { inner: OsReadRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReadRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReadRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReadRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;
    use self::libc::{c_int, size_t};

    pub struct OsRng;

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng)
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    use std::io;
    use std::mem;
    use std::ptr;
    use Rng;

    type BOOL = i32;
    type LPCSTR = *const i8;
    type DWORD = u32;
    type HCRYPTPROV = usize;
    type BYTE = u8;

    const PROV_RSA_FULL: DWORD = 1;
    const CRYPT_SILENT: DWORD = 0x00000040;
    const CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[link(name = "advapi32")]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                szContainer: LPCSTR,
                                szProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, ptr::null(), ptr::null(),
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(target_os = "nacl")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;

    pub struct OsRng(extern fn(dest: *mut libc::c_void,
                               bytes: libc::size_t,
                               read: *mut libc::size_t) -> libc::c_int);

    extern {
        fn nacl_interface_query(name: *const libc::c_char,
                                table: *mut libc::c_void,
                                table_size: libc::size_t) -> libc::size_t;
    }

    const INTERFACE: &'static [u8] = b"nacl-irt-random-0.1\0";

    #[repr(C)]
    struct NaClIRTRandom {
        get_random_bytes: Option<extern fn(dest: *mut libc::c_void,
                                           bytes: libc::size_t,
                                           read: *mut libc::size_t) -> libc::c_int>,
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            let mut iface = NaClIRTRandom {
                get_random_bytes: None,
            };
            let result = unsafe {
                nacl_interface_query(INTERFACE.as_ptr() as *const _,
                                     mem::transmute(&mut iface),
                                     mem::size_of::<NaClIRTRandom>() as libc::size_t)
            };
            if result != 0 {
                assert!(iface.get_random_bytes.is_some());
                let result = OsRng(iface.get_random_bytes.take().unwrap());
                Ok(result)
            } else {
                let error = io::ErrorKind::NotFound;
                let error = io::Error::new(error, "IRT random interface missing");
                Err(error)
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let mut read = 0;
            loop {
                let mut r: libc::size_t = 0;
                let len = v.len();
                let error = (self.0)(v[read..].as_mut_ptr() as *mut _,
                                     (len - read) as libc::size_t,
                                     &mut r as *mut _);
                assert!(error == 0, "`get_random_bytes` failed!");
                read += r as usize;

                if read >= v.len() { break; }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;
    use Rng;
    use OsRng;
    use std::thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0u8; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(()).unwrap();
        }
    }
}
<commit_msg>Fix filling buffers 4 GiB or larger with OsRng::fill_bytes on Windows<commit_after>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

use std::io;
use Rng;

/// A random number generator that retrieves randomness straight from
/// the operating system. Platform sources:
///
/// - Unix-like systems (Linux, Android, Mac OSX): read directly from
///   `/dev/urandom`, or from `getrandom(2)` system call if available.
/// - Windows: calls `CryptGenRandom`, using the default cryptographic
///   service provider with the `PROV_RSA_FULL` type.
/// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
/// - PNaCl: calls into the `nacl-irt-random-0.1` IRT interface.
///
/// This does not block.
pub struct OsRng(imp::OsRng);

impl OsRng {
    /// Create a new `OsRng`.
    pub fn new() -> io::Result<OsRng> {
        imp::OsRng::new().map(OsRng)
    }
}

impl Rng for OsRng {
    fn next_u32(&mut self) -> u32 { self.0.next_u32() }
    fn next_u64(&mut self) -> u64 { self.0.next_u64() }
    fn fill_bytes(&mut self, v: &mut [u8]) { self.0.fill_bytes(v) }
}

#[cfg(all(unix, not(target_os = "ios"),
          not(target_os = "nacl")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use std::io;
    use std::fs::File;
    use Rng;
    use read::ReadRng;
    use std::mem;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(target_arch = "aarch64")]
        const NR_GETRANDOM: libc::c_long = 278;
        #[cfg(target_arch = "powerpc")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0u8; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0u8; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn is_getrandom_available() -> bool {
        use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
        use std::sync::{Once, ONCE_INIT};

        static CHECKER: Once = ONCE_INIT;
        static AVAILABLE: AtomicBool = ATOMIC_BOOL_INIT;

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn is_getrandom_available() -> bool { false }

    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReadRng(ReadRng<File>),
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReadRng::new(reader);

            Ok(OsRng { inner: OsReadRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReadRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReadRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReadRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;
    use self::libc::{c_int, size_t};

    pub struct OsRng;

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng)
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    use std::io;
    use std::mem;
    use std::ptr;
    use Rng;

    type BOOL = i32;
    type LPCSTR = *const i8;
    type DWORD = u32;
    type HCRYPTPROV = usize;
    type BYTE = u8;

    const PROV_RSA_FULL: DWORD = 1;
    const CRYPT_SILENT: DWORD = 0x00000040;
    const CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[link(name = "advapi32")]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                szContainer: LPCSTR,
                                szProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, ptr::null(), ptr::null(),
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            // CryptGenRandom takes a DWORD (u32) for the length so we need to
            // split up the buffer.
            for slice in v.chunks_mut(<DWORD>::max_value() as usize) {
                let ret = unsafe {
                    CryptGenRandom(self.hcryptprov, slice.len() as DWORD,
                                   slice.as_mut_ptr())
                };
                if ret == 0 {
                    panic!("couldn't generate random bytes: {}",
                           io::Error::last_os_error());
                }
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(target_os = "nacl")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;

    pub struct OsRng(extern fn(dest: *mut libc::c_void,
                               bytes: libc::size_t,
                               read: *mut libc::size_t) -> libc::c_int);

    extern {
        fn nacl_interface_query(name: *const libc::c_char,
                                table: *mut libc::c_void,
                                table_size: libc::size_t) -> libc::size_t;
    }

    const INTERFACE: &'static [u8] = b"nacl-irt-random-0.1\0";

    #[repr(C)]
    struct NaClIRTRandom {
        get_random_bytes: Option<extern fn(dest: *mut libc::c_void,
                                           bytes: libc::size_t,
                                           read: *mut libc::size_t) -> libc::c_int>,
    }

    impl OsRng {
        pub fn new() -> io::Result<OsRng> {
            let mut iface = NaClIRTRandom {
                get_random_bytes: None,
            };
            let result = unsafe {
                nacl_interface_query(INTERFACE.as_ptr() as *const _,
                                     mem::transmute(&mut iface),
                                     mem::size_of::<NaClIRTRandom>() as libc::size_t)
            };
            if result != 0 {
                assert!(iface.get_random_bytes.is_some());
                let result = OsRng(iface.get_random_bytes.take().unwrap());
                Ok(result)
            } else {
                let error = io::ErrorKind::NotFound;
                let error = io::Error::new(error, "IRT random interface missing");
                Err(error)
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let mut read = 0;
            loop {
                let mut r: libc::size_t = 0;
                let len = v.len();
                let error = (self.0)(v[read..].as_mut_ptr() as *mut _,
                                     (len - read) as libc::size_t,
                                     &mut r as *mut _);
                assert!(error == 0, "`get_random_bytes` failed!");
                read += r as usize;

                if read >= v.len() { break; }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;
    use Rng;
    use OsRng;
    use std::thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0u8; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>use std::usize;

// FIXME: prefer this over Option<usize> (won't require another byte of memory)
const _MARKER: usize = usize::MAX;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SuffixType {
    Small,
    Large,
    LeftMostSmall,
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
pub fn suffix_array(input: Vec<u8>) -> Vec<usize> {
    let mut type_map = vec![SuffixType::Small; input.len() + 1];
    let mut bucket_sizes = vec![0; 256];
    let mut bucket_heads = vec![0; 256];
    let mut bucket_tails = vec![0; 256];

    type_map[input.len()] = SuffixType::LeftMostSmall;      // assume that this is a null byte
    type_map[input.len() - 1] = SuffixType::Large;          // should be L-type
    bucket_sizes[input[input.len() - 1] as usize] = 1;      // don't forget the last value!

    // group the bytes into S-type or L-type (also mark LMS types)
    for i in (0..input.len() - 1).rev() {
        bucket_sizes[input[i] as usize] += 1;

        if input[i] > input[i + 1] ||
           (input[i] == input[i + 1] && type_map[i + 1] == SuffixType::Large) {
            if type_map[i + 1] == SuffixType::Small {
                type_map[i + 1] = SuffixType::LeftMostSmall;
            }

            type_map[i] = SuffixType::Large;
        }
    }

    let mut idx = 1;
    // fill the bucket heads and tails
    for (i, size) in bucket_sizes.iter().enumerate() {
        bucket_heads[i] = idx;
        idx += *size;
        bucket_tails[i] = idx - 1;
    }

    let mut approx_sa = {       // build the first (approximate) SA
        let mut vec = vec![-1i32; input.len() + 1];
        let mut bucket_tails = bucket_tails.clone();
        for i in 0..input.len() {
            if type_map[i] != SuffixType::LeftMostSmall {
                continue        // ignore the LMS types for now
            }

            let bucket_idx = input[i] as usize;
            vec[bucket_tails[bucket_idx]] = i as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        vec[0] = input.len() as i32;       // null byte
        vec
    };

    fn induced_sort_large(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_heads: Vec<usize>, type_map: &[SuffixType]) {
        for i in 0..approx_sa.len() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] != SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_heads[bucket_idx]] = j;
            bucket_heads[bucket_idx] += 1;
        }
    }

    fn induced_sort_small(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_tails: Vec<usize>, type_map: &[SuffixType]) {
        for i in (0..approx_sa.len()).rev() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] == SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_tails[bucket_idx]] = j;
            bucket_tails[bucket_idx] -= 1;
        }
    }

    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    fn is_equal_lms(input: &[u8], type_map: &[SuffixType], j: usize, k: usize) -> bool {
        if j == input.len() || k == input.len() {
            return false    // null byte
        }

        for i in 0..(input.len() + 1) {
            let first_lms = type_map[i + j] == SuffixType::LeftMostSmall;
            let second_lms = type_map[i + k] == SuffixType::LeftMostSmall;
            if first_lms && second_lms && i > 0 {
                return true
            } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
                return false
            }
        }

        false
    }

    let mut lms_bytes = vec![input.len(); input.len() + 1];
    let mut byte = 0;
    lms_bytes[approx_sa[0] as usize] = 0;
    let mut last_idx = approx_sa[0] as usize;

    for i in 1..approx_sa.len() {
        let idx = if approx_sa[i] == -1 { input.len() } else { approx_sa[i] as usize };
        if type_map[idx] != SuffixType::LeftMostSmall {
            continue
        }

        if !is_equal_lms(&input, &type_map, last_idx, idx) {
            byte += 1;
        }

        last_idx = idx;
        lms_bytes[idx] = byte;
    }

    drop(approx_sa);

    let mut summary_string = Vec::with_capacity(lms_bytes.len());
    let mut summary_index = Vec::with_capacity(lms_bytes.len());
    for (i, b) in lms_bytes.drain(..).enumerate() {
        if b == input.len() {
            continue
        }

        summary_index.push(i);
        summary_string.push(b as u8);
    }

    let mut sa = {
        let mut bucket_tails = bucket_tails.clone();
        let summary_sa = if byte == summary_string.len() - 1 {
            let mut sa = vec![0usize; summary_string.len() + 1];
            sa[0] = summary_string.len();
            for i in 0..summary_string.len() {
                let idx = summary_string[i];
                sa[idx as usize + 1] = i;
            }

            sa
        } else {
            suffix_array(summary_string)
        };

        let mut suffix_idx = vec![-1i32; input.len() + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index[summary_sa[i] as usize];
            let bucket_idx = *input.get(idx).unwrap_or(&0) as usize;
            suffix_idx[bucket_tails[bucket_idx]] = idx as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        suffix_idx[0] = input.len() as i32;

        suffix_idx
    };

    induced_sort_large(&input, &mut sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut sa, bucket_tails, &type_map);

    sa.into_iter().map(|i| i as usize).collect()
}
<commit_msg>Prefer usize::MAX for marker<commit_after>use std::usize;

// Prefer this for marking, instead of Option<usize> (as it requires additional byte of memory)
const MARKER: usize = usize::MAX;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SuffixType {
    Small,
    Large,
    LeftMostSmall,
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
pub fn suffix_array(input: Vec<u8>) -> Vec<usize> {
    let mut type_map = vec![SuffixType::Small; input.len() + 1];
    let mut bucket_sizes = [0; 256];
    let mut bucket_heads = vec![0; 256];
    let mut bucket_tails = vec![0; 256];

    type_map[input.len()] = SuffixType::LeftMostSmall;      // assume that this is a null byte
    type_map[input.len() - 1] = SuffixType::Large;          // should be L-type
    bucket_sizes[input[input.len() - 1] as usize] = 1;      // don't forget the last value!

    // group the bytes into S-type or L-type (also mark LMS types)
    for i in (0..input.len() - 1).rev() {
        bucket_sizes[input[i] as usize] += 1;

        if input[i] > input[i + 1] ||
           (input[i] == input[i + 1] && type_map[i + 1] == SuffixType::Large) {
            if type_map[i + 1] == SuffixType::Small {
                type_map[i + 1] = SuffixType::LeftMostSmall;
            }

            type_map[i] = SuffixType::Large;
        }
    }

    let mut idx = 1;
    // fill the bucket heads and tails
    for (i, size) in bucket_sizes.iter().enumerate() {
        bucket_heads[i] = idx;
        idx += *size;
        bucket_tails[i] = idx - 1;
    }

    let mut approx_sa = {       // build the first (approximate) SA
        let mut vec = vec![MARKER; input.len() + 1];
        let mut bucket_tails = bucket_tails.clone();
        for (i, byte) in input.iter().enumerate() {
            if type_map[i] != SuffixType::LeftMostSmall {
                continue        // ignore the L and S types (for now)
            }

            let bucket_idx = *byte as usize;
            vec[bucket_tails[bucket_idx]] = i;
            bucket_tails[bucket_idx] -= 1;
        }

        vec[0] = input.len();       // null byte
        vec
    };

    fn induced_sort_large(input: &[u8], approx_sa: &mut [usize],
                          mut bucket_heads: Vec<usize>, type_map: &[SuffixType]) {
        for i in 0..approx_sa.len() {
            if approx_sa[i] == MARKER || approx_sa[i] == 0 {
                continue
            }

            let j = approx_sa[i] - 1;
            if type_map[j] != SuffixType::Large {
                continue    // only the L-types
            }

            let bucket_idx = input[j] as usize;
            approx_sa[bucket_heads[bucket_idx]] = j;
            bucket_heads[bucket_idx] += 1;
        }
    }

    fn induced_sort_small(input: &[u8], approx_sa: &mut [usize],
                          mut bucket_tails: Vec<usize>, type_map: &[SuffixType]) {
        for i in (0..approx_sa.len()).rev() {
            if approx_sa[i] == MARKER || approx_sa[i] == 0 {
                continue
            }

            let j = approx_sa[i] - 1;
            if type_map[j] == SuffixType::Large {
                continue    // only the S-types (and LMS-types as per our grouping)
            }

            let bucket_idx = input[j] as usize;
            approx_sa[bucket_tails[bucket_idx]] = j;
            bucket_tails[bucket_idx] -= 1;
        }
    }

    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    fn is_equal_lms(input: &[u8], type_map: &[SuffixType], j: usize, k: usize) -> bool {
        if j == input.len() || k == input.len() {
            return false    // null byte
        }

        for i in 0..(input.len() + 1) {
            let first_lms = type_map[i + j] == SuffixType::LeftMostSmall;
            let second_lms = type_map[i + k] == SuffixType::LeftMostSmall;
            if first_lms && second_lms && i > 0 {
                return true
            } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
                return false
            }
        }

        false
    }

    let mut byte = 0;
    let lms_bytes = {
        let mut approx_sa = approx_sa;
        let mut last_idx = approx_sa[0];
        let mut lms_bytes = vec![input.len(); input.len() + 1];
        lms_bytes[last_idx] = 0;

        for count in approx_sa.drain(1..) {
            let idx = if count == MARKER { input.len() } else { count };
            if type_map[idx] != SuffixType::LeftMostSmall {
                continue
            }

            if !is_equal_lms(&input, &type_map, last_idx, idx) {
                byte += 1;
            }

            last_idx = idx;
            lms_bytes[idx] = byte;
        }

        lms_bytes
    };

    let summary_index = lms_bytes.iter().enumerate().filter_map(|(i, b)| {
        if *b == input.len() { None } else { Some(i) }
    }).collect::<Vec<_>>();

    let mut final_sa = {
        let mut bucket_tails = bucket_tails.clone();
        let summary_sa = if byte == summary_index.len() - 1 {
            let mut sum_sa = vec![0; summary_index.len() + 1];
            sum_sa[0] = summary_index.len();
            for i in 0..summary_index.len() {
                let idx = lms_bytes[summary_index[i]] as usize;
                sum_sa[idx + 1] = i;
            }

            sum_sa
        } else {
            suffix_array(summary_index.iter().map(|i| lms_bytes[*i] as u8).collect())
        };

        let mut suffix_idx = vec![MARKER; input.len() + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index[summary_sa[i]];
            let bucket_idx = input[idx] as usize;
            suffix_idx[bucket_tails[bucket_idx]] = idx;
            bucket_tails[bucket_idx] -= 1;
        }

        suffix_idx[0] = input.len();
        suffix_idx
    };

    induced_sort_large(&input, &mut final_sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut final_sa, bucket_tails, &type_map);

    final_sa
}
<|endoftext|>