<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::println;
use std::cell::Cell;
use std::{os, str, io, run, hashmap, task};
use extra::arc;
use extra::priority_queue::PriorityQueue;
use std::comm::*;
use std::cast;
use extra::comm::DuplexStream;
use std::path::Path;
use std::num;
use extra::sort;

static PORT:  int = 4414;
static IP: &'static str = "127.0.0.1";
static mut visitor_count: uint = 0;

#[deriving(Clone)]
struct access_t {
    filepath: ~std::path::PosixPath,
    size: i64,
    num_access: uint
}

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    priority: uint
}

impl std::cmp::Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {
        return self.priority > other.priority;
    }
}

struct IpRange {
    start_addr: u32,
    end_addr: u32
}

fn main() {

    let req_heap: PriorityQueue<sched_msg> = PriorityQueue::new();
    let shared_req_heap = arc::RWArc::new(req_heap);
    let add_vec = shared_req_heap.clone();
    let take_vec = shared_req_heap.clone();

    let cache: hashmap::HashMap<~str, ~str> = hashmap::HashMap::new();
    let shared_cache = arc::RWArc::new(cache);
    let add_cache = shared_cache.clone();
    let check_cache = shared_cache.clone();

    let accesses: ~[access_t] = ~[];
    let shared_accesses = arc::RWArc::new(accesses);
    let exist_accesses = shared_accesses.clone();
    let new_accesses = shared_accesses.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                let fpath = tf.filepath.clone();
                println(fmt!("begin serving file [%?]", tf.filepath.to_str()));
                // A web server should always reply a HTTP header for any legal HTTP request.
                tf.stream.write("HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n".as_bytes());
                match do check_cache.read |cc| { (*cc).find_copy(&tf.filepath.to_str()) } {
                    None        =>  { if tf.filepath.get_mode().unwrap() % 2 == 1 {
                                            execFile(tf);
                                            // No cache for dynamically-generated files.
                                            loop;
                                      }
                                      let data = writeFile(&mut tf);
                                      do new_accesses.write |na| {
                                            if (*na).len() > 20 {
                                                (*na).pop_opt();
                                            }

                                            let file_info = access_t{filepath: fpath.clone(), size: fpath.stat().unwrap().st_size, num_access: 1};
                                            (*na).push(file_info.clone());
                                            do add_cache.write |ac| {
                                                let sort_access = sort::merge_sort(*na, |it1: &access_t, it2: &access_t| {
                                                                                (it1.size*it1.num_access as i64) <= (it2.size*it2.num_access as i64)
                                                                                });
                                                for item in sort_access.iter() {
                                                    if file_info.size > (item.size*item.num_access as i64) {
                                                        match (*ac).pop(&item.filepath.to_str()) {
                                                            Some(_)    =>  { (*ac).swap(file_info.filepath.to_str(), data.clone()); break;}
                                                            None           =>  { }
                                                        }
                                                    }
                                                }
                                            }
                                      }
                                  }

                    Some(ct)    =>  { tf.stream.write(ct.as_bytes());
                                      do exist_accesses.write |ea| {
                                            for i in range(0, (*ea).len()) {
                                                let mut item = (*ea)[i].clone();
                                                if item.filepath == tf.filepath {
                                                    item.num_access += 1;
                                                    (*ea).remove(i);
                                                    (*ea).push(item);
                                                    break;
                                                }
                                            }
                                        }
                                    }
                }
                println(fmt!("finish file [%?]", fpath.to_str()));
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {
                    // LIFO didn't make sense in service scheduling, so we modify it as FIFO by using shift_opt() rather than pop().
                    let tf: sched_msg = (*vec).pop();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    // turns file into list of strings for each line
    let iplist: ~str = io::read_whole_file_str(&Path("iplist.txt")).unwrap();
    let iplist: ~[~[u32]] = iplist.line_iter().map(|x| {
                                                            let y: ~[&str] = x.split_iter(' ').collect();
                                                            let start: u32 = num::from_str_radix(y[0], 10).unwrap();
                                                            let fin: u32 = num::from_str_radix(y[1], 10).unwrap();
                                                            ~[start, fin]
                                                         }).collect();
    let iplist = sort::merge_sort(iplist, |x, y| {
                                                     x[0] <= y[0]
                                                    });

    //println(fmt!("Checking! %?", check_ip(1280123655, iplist)));
    let shared_ip_map = arc::RWArc::new(iplist);

    let shared_count = arc::RWArc::new(0);

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    let mut acceptor = socket.listen().unwrap();

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    for stream in acceptor.incoming() {
        let stream = Cell::new(stream);

        let incr_count = shared_count.clone();

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let shared_ip_map = shared_ip_map.clone();
        let child_add_vec = add_vec.clone();

        do spawn {
            do incr_count.write |count| {
                *count += 1;
            }

            let mut stream = stream.take();
            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));
                // More better path security!
                let unclean_path = os::getcwd().push(Path(path).to_str()).to_str();
                let mut file_path = ~os::getcwd();
                // paths are always normalized so a/b/../c becomes a/c
                if unclean_path.starts_with(file_path.to_str()) {
                    file_path = ~file_path.push(path);
                }
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", incr_count.read(|c| { *c }));

                    stream.write(response.as_bytes());
                }
                else {
                    // Requests scheduling

                    let mut priority = file_path.stat().unwrap().st_size as uint;
                    unsafe {
                        match stream {
                            Some(ref s) => {
                                    let stream = cast::transmute_mut(s);
                                    let pn = stream.peer_name().unwrap();
                                    println(fmt!("Peer is: %?", pn));
                                    match pn.ip {
                                        Ipv4Addr(a, b, c, d) => {
                                                                // Since we are sharing the ip_map it must be read
                                                                do shared_ip_map.read |map| {
                                                                    if fcheck_ip(a,b,c,d, *map) {
                                                                        priority = 1;
                                                                        println("local request!");
                                                                    }
                                                                }
                                                            },
                                        _                    =>  fail!()
                                    }
                            },
                            _    => fail!()
                        };
                    }
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone(), priority: priority};
                    let (sm_port, sm_chan) = std::comm::stream();
                    sm_chan.send(msg);

                    do child_add_vec.write |vec| {
                        let msg = sm_port.recv();
                        (*vec).push(msg); // enqueue new request.
                        println("add to queue");
                    }
                    child_chan.send(""); //notify the new arriving request.
                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}

fn fcheck_ip(a: u8, b: u8, c: u8, d: u8, iplist: &[~[u32]]) -> bool {
    check_ip(a as u32 << 24 + b as u32 << 16 + c as u8 << 8 + d, iplist)
}

// Looks up an ip prefix in the hashset by trying each octet
fn check_ip(ip: u32, iplist: &[~[u32]]) -> bool {
    let mut min = 0;
    let mut max = iplist.len() - 1;
    while max >= min {
        let mid = (max - min)/2 + min;
        let elem = &iplist[mid];
        if ip >= elem[0] {
            if ip <= elem[1] {
                return true;
            } else {
                min = mid + 1;
            }
        } else { // ip < elem
           if max == 0 { // because overflow thats why
                return false;
           }
           max = mid - 1;
        }
    }
    return false
}

fn execFile(file_data: sched_msg) {
    let (port, chan) = DuplexStream();
    do task::spawn_supervised {
        do_gash(&chan)
    }
    let mut file_data = file_data;
    match io::file_reader(file_data.filepath) {
        Ok(rd)      =>  {   let closer = ~['\"',' ', '-','-','>'];
                            while !rd.eof() {
                            let rd_byte = rd.read_byte() as u8;

                            match rd_byte {
                                0x3c =>  { let mut open: ~[u8] = ~[0,0,0,0,0,0,0,0,0,0,0,0,0,0];
                                                          rd.read(open, 14);
                                                          if str::from_utf8(open) == ~"!--#exec cmd=\"" {
                                                                let mut cmd: ~[u8] = ~[];
                                                                let mut cmd_byte: u8;
                                                                let mut i = 0;
                                                                loop {
                                                                    cmd_byte = rd.read_byte() as u8;
                                                                    if cmd_byte == closer[i] as u8 {
                                                                        if i == closer.len() - 1 {
                                                                            break;
                                                                        }

                                                                        i += 1;
                                                                    }
                                                                    else {
                                                                        i = 0;
                                                                    }

                                                                    if cmd_byte != 0xFF {
                                                                        cmd.push(cmd_byte);
                                                                    }
                                                                }

                                                          port.send(str::from_utf8(cmd.slice_to(cmd.len() - 4)));
                                                          let result = port.recv();
                                                          file_data.stream.write(result.as_bytes());
                                                          }
                                                        },
                                _                   =>  { if rd_byte != 0xFF {
                                                                file_data.stream.write(&[rd_byte]);
                                                          }
                                                        }
                            }
                          }
                          port.send(~"end");
                        }
        Err(err)    =>  { println(err); }
    }

}

fn do_gash(chan: &DuplexStream<~str, ~str>) {
    let mut gash = run::Process::new("./gash", &[], run::ProcessOptions::new());
    let mut cmd: ~str;
    let mut result: ~str;
    let gin = gash.input();
    let gout = gash.output();
    let mut res: ~[u8];
    let mut prompt: ~[u8] = ~[];
    let mut res_byte: u8;
    for _ in range(0,8) {
        prompt.push(0);
    }

    gout.read(prompt, 8);
    // Insecure. We should use some other means of communicating the need to die.
    loop {
        cmd = chan.recv();
        if cmd == ~"end" {
            break;
        }

        cmd.push_str("\n");
        res = ~[];
        gin.write(cmd.as_bytes());
        loop {
            res_byte = gout.read_byte() as u8;
            if res_byte == '\0' as u8{
                break;
            }

            res.push(res_byte);
        }

        gout.read(prompt, 7);
        result = str::from_utf8(res);
        chan.send(result.clone());
    }

    gash.destroy();
}

fn writeFile(tf: &mut sched_msg) ->  ~str {
    let mut file: ~[u8] = ~[];
    let mut result: ~str;
    match io::file_reader(tf.filepath) {
        Ok(rd)      =>  { while !rd.eof() {
                            let rd_byte = rd.read_byte() as u8;
                            if rd_byte != 0xFF {
                                tf.stream.write(&[rd_byte]);
                                file.push(rd_byte);
                            }
                          }
                            result = str::from_utf8(file);
                        }
        Err(err)    =>  { println(err); result = err; }
    }
    result
}

<commit_msg>premerge<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::println;
use std::cell::Cell;
use std::{os, str, io, run, hashmap, task};
use extra::arc;
use extra::priority_queue::PriorityQueue;
use std::comm::*;
use std::cast;
use extra::comm::DuplexStream;
use std::path::Path;
use std::num;
use extra::{sort, json};

static PORT:  int = 4414;
static IP: &'static str = "127.0.0.1";
static mut visitor_count: uint = 0;

#[deriving(Clone)]
struct access_t {
    filepath: ~std::path::PosixPath,
    size: i64,
    num_access: uint
}

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    priority: uint
}

impl std::cmp::Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {
        return self.priority > other.priority;
    }
}

fn main() {
    let s: ~str = io::read_whole_file_str(&Path("server.conf")).unwrap();
    let js: json::Json = json::from_str(s).unwrap();
    println(fmt!("%?", js));

    let req_heap: PriorityQueue<sched_msg> = PriorityQueue::new();
    let shared_req_heap = arc::RWArc::new(req_heap);
    let add_vec = shared_req_heap.clone();
    let take_vec = shared_req_heap.clone();

    let cache: hashmap::HashMap<~str, ~str> = hashmap::HashMap::new();
    let shared_cache = arc::RWArc::new(cache);
    let add_cache = shared_cache.clone();
    let check_cache = shared_cache.clone();

    let accesses: ~[access_t] = ~[];
    let shared_accesses = arc::RWArc::new(accesses);
    let exist_accesses = shared_accesses.clone();
    let new_accesses = shared_accesses.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle
                let fpath = tf.filepath.clone();
                println(fmt!("begin serving file [%?]", tf.filepath.to_str()));
                // A web server should always reply a HTTP header for any legal HTTP request.
                tf.stream.write("HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n".as_bytes());
                match do check_cache.read |cc| { (*cc).find_copy(&tf.filepath.to_str()) } {
                    None        =>  { if tf.filepath.get_mode().unwrap() % 2 == 1 {
                                            execFile(tf);
                                            // No cache for dynamically-generated files.
                                            loop;
                                      }
                                      let data = writeFile(&mut tf);
                                      do new_accesses.write |na| {
                                            if (*na).len() > 20 {
                                                (*na).pop_opt();
                                            }

                                            let file_info = access_t{filepath: fpath.clone(), size: fpath.stat().unwrap().st_size, num_access: 1};
                                            (*na).push(file_info.clone());
                                            do add_cache.write |ac| {
                                                let sort_access = sort::merge_sort(*na, |it1: &access_t, it2: &access_t| {
                                                                                (it1.size*it1.num_access as i64) <= (it2.size*it2.num_access as i64)
                                                                                });
                                                for item in sort_access.iter() {
                                                    if file_info.size > (item.size*item.num_access as i64) {
                                                        match (*ac).pop(&item.filepath.to_str()) {
                                                            Some(_)    =>  { (*ac).swap(file_info.filepath.to_str(), data.clone()); break;}
                                                            None           =>  { }
                                                        }
                                                    }
                                                }
                                            }
                                      }
                                  }

                    Some(ct)    =>  { tf.stream.write(ct.as_bytes());
                                      do exist_accesses.write |ea| {
                                            for i in range(0, (*ea).len()) {
                                                let mut item = (*ea)[i].clone();
                                                if item.filepath == tf.filepath {
                                                    item.num_access += 1;
                                                    (*ea).remove(i);
                                                    (*ea).push(item);
                                                    break;
                                                }
                                            }
                                        }
                                    }
                }
                println(fmt!("finish file [%?]", fpath.to_str()));
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {
                    // LIFO didn't make sense in service scheduling, so we modify it as FIFO by using shift_opt() rather than pop().
                    let tf: sched_msg = (*vec).pop();
                    println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    // turns file into list of strings for each line
    let iplist: ~str = io::read_whole_file_str(&Path("iplist.txt")).unwrap();
    let iplist: ~[~[u32]] = iplist.line_iter().map(|x| {
                                                            let y: ~[&str] = x.split_iter(' ').collect();
                                                            let start: u32 = num::from_str_radix(y[0], 10).unwrap();
                                                            let fin: u32 = num::from_str_radix(y[1], 10).unwrap();
                                                            ~[start, fin]
                                                         }).collect();
    let iplist = sort::merge_sort(iplist, |x, y| {
                                                     x[0] <= y[0]
                                                    });

    //println(fmt!("Checking! %?", check_ip(1280123655, iplist)));
    let shared_ip_map = arc::RWArc::new(iplist);

    let shared_count = arc::RWArc::new(0);

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    let mut acceptor = socket.listen().unwrap();

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    for stream in acceptor.incoming() {
        let stream = Cell::new(stream);

        let incr_count = shared_count.clone();

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let shared_ip_map = shared_ip_map.clone();
        let child_add_vec = add_vec.clone();

        do spawn {
            do incr_count.write |count| {
                *count += 1;
            }

            let mut stream = stream.take();
            let mut buf = [0, ..500];
            stream.read(buf);
            let request_str = str::from_utf8(buf);

            let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
            if req_group.len() > 2 {
                let path = req_group[1];
                println(fmt!("Request for path: \n%?", path));
                // More better path security!
                let unclean_path = os::getcwd().push(Path(path).to_str()).to_str();
                let mut file_path = ~os::getcwd();
                // paths are always normalized so a/b/../c becomes a/c
                if unclean_path.starts_with(file_path.to_str()) {
                    file_path = ~file_path.push(path);
                }
                if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                    println(fmt!("Request received:\n%s", request_str));
                    let response: ~str = fmt!(
                        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                         <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                         \r\n", incr_count.read(|c| { *c }));

                    stream.write(response.as_bytes());
                }
                else {
                    // Requests scheduling

                    let mut priority = file_path.stat().unwrap().st_size as uint;
                    unsafe {
                        match stream {
                            Some(ref s) => {
                                    let stream = cast::transmute_mut(s);
                                    let pn = stream.peer_name().unwrap();
                                    println(fmt!("Peer is: %?", pn));
                                    match pn.ip {
                                        Ipv4Addr(a, b, c, d) => {
                                                                // Since we are sharing the ip_map it must be read
                                                                do shared_ip_map.read |map| {
                                                                    if fcheck_ip(a,b,c,d, *map) {
                                                                        priority = 1;
                                                                        println("local request!");
                                                                    }
                                                                }
                                                            },
                                        _                    =>  fail!()
                                    }
                            },
                            _    => fail!()
                        };
                    }
                    let msg: sched_msg = sched_msg{stream: stream, filepath: file_path.clone(), priority: priority};
                    let (sm_port, sm_chan) = std::comm::stream();
                    sm_chan.send(msg);

                    do child_add_vec.write |vec| {
                        let msg = sm_port.recv();
                        (*vec).push(msg); // enqueue new request.
                        println("add to queue");
                    }
                    child_chan.send(""); //notify the new arriving request.
                    println(fmt!("get file request: %?", file_path));
                }
            }
            println!("connection terminates")
        }
    }
}

fn fcheck_ip(a: u8, b: u8, c: u8, d: u8, iplist: &[~[u32]]) -> bool {
    check_ip(a as u32 << 24 + b as u32 << 16 + c as u8 << 8 + d, iplist)
}

// Looks up an ip prefix in the hashset by trying each octet
fn check_ip(ip: u32, iplist: &[~[u32]]) -> bool {
    let mut min = 0;
    let mut max = iplist.len() - 1;
    while max >= min {
        let mid = (max - min)/2 + min;
        let elem = &iplist[mid];
        if ip >= elem[0] {
            if ip <= elem[1] {
                return true;
            } else {
                min = mid + 1;
            }
        } else { // ip < elem
           if max == 0 { // because overflow thats why
                return false;
           }
           max = mid - 1;
        }
    }
    return false
}

fn execFile(file_data: sched_msg) {
    let (port, chan) = DuplexStream();
    do task::spawn_supervised {
        do_gash(&chan)
    }
    let mut file_data = file_data;
    match io::file_reader(file_data.filepath) {
        Ok(rd)      =>  {   let closer = ~['\"',' ', '-','-','>'];
                            while !rd.eof() {
                            let rd_byte = rd.read_byte() as u8;

                            match rd_byte {
                                0x3c =>  { let mut open: ~[u8] = ~[0,0,0,0,0,0,0,0,0,0,0,0,0,0];
                                                          rd.read(open, 14);
                                                          if str::from_utf8(open) == ~"!--#exec cmd=\"" {
                                                                let mut cmd: ~[u8] = ~[];
                                                                let mut cmd_byte: u8;
                                                                let mut i = 0;
                                                                loop {
                                                                    cmd_byte = rd.read_byte() as u8;
                                                                    if cmd_byte == closer[i] as u8 {
                                                                        if i == closer.len() - 1 {
                                                                            break;
                                                                        }

                                                                        i += 1;
                                                                    }
                                                                    else {
                                                                        i = 0;
                                                                    }

                                                                    if cmd_byte != 0xFF {
                                                                        cmd.push(cmd_byte);
                                                                    }
                                                                }

                                                          port.send(str::from_utf8(cmd.slice_to(cmd.len() - 4)));
                                                          let result = port.recv();
                                                          file_data.stream.write(result.as_bytes());
                                                          }
                                                        },
                                _                   =>  { if rd_byte != 0xFF {
                                                                file_data.stream.write(&[rd_byte]);
                                                          }
                                                        }
                            }
                          }
                          port.send(~"end");
                        }
        Err(err)    =>  { println(err); }
    }

}

fn do_gash(chan: &DuplexStream<~str, ~str>) {
    let mut gash = run::Process::new("./gash", &[], run::ProcessOptions::new());
    let mut cmd: ~str;
    let mut result: ~str;
    let gin = gash.input();
    let gout = gash.output();
    let mut res: ~[u8];
    let mut prompt: ~[u8] = ~[];
    let mut res_byte: u8;
    for _ in range(0,8) {
        prompt.push(0);
    }

    gout.read(prompt, 8);
    // Insecure. We should use some other means of communicating the need to die.
    loop {
        cmd = chan.recv();
        if cmd == ~"end" {
            break;
        }

        cmd.push_str("\n");
        res = ~[];
        gin.write(cmd.as_bytes());
        loop {
            res_byte = gout.read_byte() as u8;
            if res_byte == '\0' as u8{
                break;
            }

            res.push(res_byte);
        }

        gout.read(prompt, 7);
        result = str::from_utf8(res);
        chan.send(result.clone());
    }

    gash.destroy();
}

fn writeFile(tf: &mut sched_msg) ->  ~str {
    let mut file: ~[u8] = ~[];
    let mut result: ~str;
    match io::file_reader(tf.filepath) {
        Ok(rd)      =>  { while !rd.eof() {
                            let rd_byte = rd.read_byte() as u8;
                            if rd_byte != 0xFF {
                                tf.stream.write(&[rd_byte]);
                                file.push(rd_byte);
                            }
                          }
                            result = str::from_utf8(file);
                        }
        Err(err)    =>  { println(err); result = err; }
    }
    result
}

<|endoftext|>


<commit_before>use ast;
use ast::Ast;
use ast::Function;
use ast::Elem::ElemFunction;

use error::ErrorCode;
use error::ParseError;

use interner::Interner;

use parser::Parser;

use sym;
use sym::BuiltinType;
use sym::SymTable;
use sym::Sym::{self, SymFunction};
use sym::SymFunctionType;
use sym::Type;

macro_rules! err {
    ( $errs: expr, $x: expr ) => {
        {
            match $x {
                Ok(val) => val,
                Err(_) => { return Err($errs); }
            }
        }
    };
}

// Only do semantic analysis until some amount of
// errors found
static MAX_ERRORS: usize = 5;

pub struct SemCheck<'a> {
    ast: &'a Ast,
    interner: &'a mut Interner,
    errors: Vec<ParseError>,
}

impl<'a> SemCheck<'a> {
    pub fn new<'b>(ast: &'b Ast, interner: &'b mut Interner) -> SemCheck<'b> {
        SemCheck {
            ast: ast,
            interner: interner,
            errors: Vec::new(),
        }
    }

    pub fn check(mut self) -> Result<SymTable, Vec<ParseError>> {
        let mut globals = SymTable::new();

        add_predefined_types(self.interner, &mut globals);
        add_predefined_functions(self.interner, &mut globals);

        err!(self.errors, self.parse_function_headers(&mut globals));
        err!(self.errors, self.parse_function_bodies(&mut globals));

        if self.errors.is_empty() {
            Ok(globals)
        } else {
            Err(self.errors)
        }
    }

    fn parse_function_headers(&mut self, globals: &mut SymTable) -> Result<(), ()> {
        for elem in &self.ast.elements {
            match *elem {
                ElemFunction(ref fct) => {
                    try!(self.add_function_header(globals, fct));
                }

                _ => unreachable!()
            }
        }

        Ok(())
    }

    fn parse_function_bodies(&mut self, globals: &mut SymTable) -> Result<(), ()> {
        for fct in globals.functions_mut() {
            try!(self.check_function_body(fct));
        }

        Ok(())
    }

    fn parse_type(&mut self, globals: &SymTable, ty: &ast::Type) -> Result<sym::Type, ()> {
        let newtype = match *ty {
            ast::Type::TypeBasic(ref val) => {
                if let None = globals.find(val.name).map(|s| s.to_type()) {
                    let message = format!("unknown type {}", self.interner.str(val.name));

                    try!(self.error(ParseError {
                        code: ErrorCode::UnknownType,
                        message: message,
                        position: ty.pos()
                    }));

                    return Err(());
                }

                sym::Type::create_basic(val.name)
            }

            ast::Type::TypeTuple(ref val) => {
                let mut subtypes = Vec::with_capacity(val.subtypes.len());

                for ty in &val.subtypes {
                    let ty = box try!(self.parse_type(globals, ty));
                    subtypes.push(ty);
                }

                sym::Type::create_tuple(subtypes)
            }

            ast::Type::TypePtr(ref val) => {
                let subtype = box try!(self.parse_type(globals, &val.subtype));

                sym::Type::create_ptr(subtype)
            }

            ast::Type::TypeArray(ref val) => {
                let subtype = box try!(self.parse_type(globals, &val.subtype));

                sym::Type::create_array(subtype)
            }
        };

        Ok(newtype)
    }

    fn add_function_header(&mut self, globals: &mut SymTable, fct: &Function) -> Result<(), ()> {
        let params = try!(self.parse_function_params(globals, fct));

        let ty = if let Some(ref ty) = fct.return_type {
            try!(self.parse_type(globals, ty))
        } else {
            Type::create_unit()
        };

        let symfct = SymFunction(SymFunctionType {
            name: fct.name,
            return_type: ty,
            params: params,
            body: fct.block.id()
        });

        if let Err(_) = globals.insert(fct.name, symfct) {
            let message = format!("identifier {} already exists", self.interner.str(fct.name));

            try!(self.error(ParseError {
                code: ErrorCode::IdentifierAlreadyExists,
                message: message,
                position: fct.pos
            }));
        }

        Ok(())
    }

    fn parse_function_params(&mut self, globals: &mut SymTable, fct: &Function)
            -> Result<Vec<sym::Param>, ()> {
        let mut params = Vec::with_capacity(fct.params.len());

        for param in &fct.params {
            let p = try!(self.parse_param(globals, param));
            params.push(p);
        }

        Ok(params)
    }

    fn parse_param(&mut self, globals: &SymTable, param: &ast::Param)
            -> Result<sym::Param, ()> {
        let ty = try!(self.parse_type(globals, &param.data_type));

        Ok(sym::Param {
            name: param.name,
            data_type: ty
        })
    }

    fn check_function_body(&mut self, fct: &SymFunctionType) -> Result<(), ()> {
        Ok(())
    }

    fn error(&mut self, error: ParseError) -> Result<(), ()> {
        self.errors.push(error);

        if self.errors.len() <= MAX_ERRORS {
            Ok(())
        } else {
            Err(())
        }
    }

    fn check_function(&mut self, fct: &Function) {
        println!("check_function");
    }
}

fn add_predefined_types(interner: &mut Interner, globals: &mut SymTable) {
    let name = interner.intern(String::from("int"));
    let sym = Sym::create_type(name, BuiltinType::Int);
    globals.insert(name, sym);

    let name = interner.intern(String::from("str"));
    let sym = Sym::create_type(name, BuiltinType::Str);
    globals.insert(name, sym);

    let name = interner.intern(String::from("bool"));
    let sym = Sym::create_type(name, BuiltinType::Bool);
    globals.insert(name, sym);
}

fn add_predefined_functions(interner: &mut Interner, globals: &mut SymTable) {
    // TODO: add print(str), print_int(int)
}

#[test]
fn test_empty_file() {
    let (prog, mut interner) = Parser::from_str("").parse().unwrap();

    SemCheck::new(&prog, &mut interner).check().unwrap();
}

#[test]
fn test_function_multiple_times() {
    let (prog, mut interner) = Parser::from_str("fn main() {} fn main() {}").parse().unwrap();
    let errors = SemCheck::new(&prog, &mut interner).check().unwrap_err();

    assert_eq!(1, errors.len());
    assert_eq!(ErrorCode::IdentifierAlreadyExists, errors[0].code);
}

#[test]
fn test_builtin_types() {
    let (prog, mut interner) = Parser::from_str("fn main(a: int, b: str) { }").parse().unwrap();
    SemCheck::new(&prog, &mut interner).check().unwrap();
}

#[test]
fn test_unknown_type() {
    let (prog, mut interner) = Parser::from_str("fn main(a: abc) { }").parse().unwrap();
    let errors = SemCheck::new(&prog, &mut interner).check().unwrap_err();

    assert_eq!(1, errors.len());
    assert_eq!(ErrorCode::UnknownType, errors[0].code);
}

<commit_msg>semck: also test bool<commit_after>use ast;
use ast::Ast;
use ast::Function;
use ast::Elem::ElemFunction;

use error::ErrorCode;
use error::ParseError;

use interner::Interner;

use parser::Parser;

use sym;
use sym::BuiltinType;
use sym::SymTable;
use sym::Sym::{self, SymFunction};
use sym::SymFunctionType;
use sym::Type;

macro_rules! err {
    ( $errs: expr, $x: expr ) => {
        {
            match $x {
                Ok(val) => val,
                Err(_) => { return Err($errs); }
            }
        }
    };
}

// Only do semantic analysis until some amount of
// errors found
static MAX_ERRORS: usize = 5;

pub struct SemCheck<'a> {
    ast: &'a Ast,
    interner: &'a mut Interner,
    errors: Vec<ParseError>,
}

impl<'a> SemCheck<'a> {
    pub fn new<'b>(ast: &'b Ast, interner: &'b mut Interner) -> SemCheck<'b> {
        SemCheck {
            ast: ast,
            interner: interner,
            errors: Vec::new(),
        }
    }

    pub fn check(mut self) -> Result<SymTable, Vec<ParseError>> {
        let mut globals = SymTable::new();

        add_predefined_types(self.interner, &mut globals);
        add_predefined_functions(self.interner, &mut globals);

        err!(self.errors, self.parse_function_headers(&mut globals));
        err!(self.errors, self.parse_function_bodies(&mut globals));

        if self.errors.is_empty() {
            Ok(globals)
        } else {
            Err(self.errors)
        }
    }

    fn parse_function_headers(&mut self, globals: &mut SymTable) -> Result<(), ()> {
        for elem in &self.ast.elements {
            match *elem {
                ElemFunction(ref fct) => {
                    try!(self.add_function_header(globals, fct));
                }

                _ => unreachable!()
            }
        }

        Ok(())
    }

    fn parse_function_bodies(&mut self, globals: &mut SymTable) -> Result<(), ()> {
        for fct in globals.functions_mut() {
            try!(self.check_function_body(fct));
        }

        Ok(())
    }

    fn parse_type(&mut self, globals: &SymTable, ty: &ast::Type) -> Result<sym::Type, ()> {
        let newtype = match *ty {
            ast::Type::TypeBasic(ref val) => {
                if let None = globals.find(val.name).map(|s| s.to_type()) {
                    let message = format!("unknown type {}", self.interner.str(val.name));

                    try!(self.error(ParseError {
                        code: ErrorCode::UnknownType,
                        message: message,
                        position: ty.pos()
                    }));

                    return Err(());
                }

                sym::Type::create_basic(val.name)
            }

            ast::Type::TypeTuple(ref val) => {
                let mut subtypes = Vec::with_capacity(val.subtypes.len());

                for ty in &val.subtypes {
                    let ty = box try!(self.parse_type(globals, ty));
                    subtypes.push(ty);
                }

                sym::Type::create_tuple(subtypes)
            }

            ast::Type::TypePtr(ref val) => {
                let subtype = box try!(self.parse_type(globals, &val.subtype));

                sym::Type::create_ptr(subtype)
            }

            ast::Type::TypeArray(ref val) => {
                let subtype = box try!(self.parse_type(globals, &val.subtype));

                sym::Type::create_array(subtype)
            }
        };

        Ok(newtype)
    }

    fn add_function_header(&mut self, globals: &mut SymTable, fct: &Function) -> Result<(), ()> {
        let params = try!(self.parse_function_params(globals, fct));

        let ty = if let Some(ref ty) = fct.return_type {
            try!(self.parse_type(globals, ty))
        } else {
            Type::create_unit()
        };

        let symfct = SymFunction(SymFunctionType {
            name: fct.name,
            return_type: ty,
            params: params,
            body: fct.block.id()
        });

        if let Err(_) = globals.insert(fct.name, symfct) {
            let message = format!("identifier {} already exists", self.interner.str(fct.name));

            try!(self.error(ParseError {
                code: ErrorCode::IdentifierAlreadyExists,
                message: message,
                position: fct.pos
            }));
        }

        Ok(())
    }

    fn parse_function_params(&mut self, globals: &mut SymTable, fct: &Function)
            -> Result<Vec<sym::Param>, ()> {
        let mut params = Vec::with_capacity(fct.params.len());

        for param in &fct.params {
            let p = try!(self.parse_param(globals, param));
            params.push(p);
        }

        Ok(params)
    }

    fn parse_param(&mut self, globals: &SymTable, param: &ast::Param)
            -> Result<sym::Param, ()> {
        let ty = try!(self.parse_type(globals, &param.data_type));

        Ok(sym::Param {
            name: param.name,
            data_type: ty
        })
    }

    fn check_function_body(&mut self, fct: &SymFunctionType) -> Result<(), ()> {
        Ok(())
    }

    fn error(&mut self, error: ParseError) -> Result<(), ()> {
        self.errors.push(error);

        if self.errors.len() <= MAX_ERRORS {
            Ok(())
        } else {
            Err(())
        }
    }

    fn check_function(&mut self, fct: &Function) {
        println!("check_function");
    }
}

fn add_predefined_types(interner: &mut Interner, globals: &mut SymTable) {
    let name = interner.intern(String::from("int"));
    let sym = Sym::create_type(name, BuiltinType::Int);
    globals.insert(name, sym);

    let name = interner.intern(String::from("str"));
    let sym = Sym::create_type(name, BuiltinType::Str);
    globals.insert(name, sym);

    let name = interner.intern(String::from("bool"));
    let sym = Sym::create_type(name, BuiltinType::Bool);
    globals.insert(name, sym);
}

fn add_predefined_functions(interner: &mut Interner, globals: &mut SymTable) {
    // TODO: add print(str), print_int(int)
}

#[test]
fn test_empty_file() {
    let (prog, mut interner) = Parser::from_str("").parse().unwrap();

    SemCheck::new(&prog, &mut interner).check().unwrap();
}

#[test]
fn test_function_multiple_times() {
    let (prog, mut interner) = Parser::from_str("fn main() {} fn main() {}").parse().unwrap();
    let errors = SemCheck::new(&prog, &mut interner).check().unwrap_err();

    assert_eq!(1, errors.len());
    assert_eq!(ErrorCode::IdentifierAlreadyExists, errors[0].code);
}

#[test]
fn test_builtin_types() {
    let (prog, mut interner) = Parser::from_str(
        "fn main(a: int, b: str, c: bool) { }").parse().unwrap();
    SemCheck::new(&prog, &mut interner).check().unwrap();
}

#[test]
fn test_unknown_type() {
    let (prog, mut interner) = Parser::from_str("fn main(a: abc) { }").parse().unwrap();
    let errors = SemCheck::new(&prog, &mut interner).check().unwrap_err();

    assert_eq!(1, errors.len());
    assert_eq!(ErrorCode::UnknownType, errors[0].code);
}

<|endoftext|>


<commit_before>/// Conversion from an `Iterator`.
///
/// By implementing `FromIterator` for a type, you define how it will be
/// created from an iterator. This is common for types which describe a
/// collection of some kind.
///
/// `FromIterator`'s [`from_iter`] is rarely called explicitly, and is instead
/// used through [`Iterator`]'s [`collect`] method. See [`collect`]'s
/// documentation for more examples.
///
/// [`from_iter`]: #tymethod.from_iter
/// [`Iterator`]: trait.Iterator.html
/// [`collect`]: trait.Iterator.html#method.collect
///
/// See also: [`IntoIterator`].
///
/// [`IntoIterator`]: trait.IntoIterator.html
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use std::iter::FromIterator;
///
/// let five_fives = std::iter::repeat(5).take(5);
///
/// let v = Vec::from_iter(five_fives);
///
/// assert_eq!(v, vec![5, 5, 5, 5, 5]);
/// ```
///
/// Using [`collect`] to implicitly use `FromIterator`:
///
/// ```
/// let five_fives = std::iter::repeat(5).take(5);
///
/// let v: Vec<i32> = five_fives.collect();
///
/// assert_eq!(v, vec![5, 5, 5, 5, 5]);
/// ```
///
/// Implementing `FromIterator` for your type:
///
/// ```
/// use std::iter::FromIterator;
///
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // and we'll implement FromIterator
/// impl FromIterator<i32> for MyCollection {
///     fn from_iter<I: IntoIterator<Item=i32>>(iter: I) -> Self {
///         let mut c = MyCollection::new();
///
///         for i in iter {
///             c.add(i);
///         }
///
///         c
///     }
/// }
///
/// // Now we can make a new iterator...
/// let iter = (0..5).into_iter();
///
/// // ... and make a MyCollection out of it
/// let c = MyCollection::from_iter(iter);
///
/// assert_eq!(c.0, vec![0, 1, 2, 3, 4]);
///
/// // collect works too!
///
/// let iter = (0..5).into_iter();
/// let c: MyCollection = iter.collect();
///
/// assert_eq!(c.0, vec![0, 1, 2, 3, 4]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
#[rustc_on_unimplemented(
    message="a collection of type `{Self}` cannot be built from an iterator \
             over elements of type `{A}`",
    label="a collection of type `{Self}` cannot be built from `std::iter::Iterator<Item={A}>`",
)]
pub trait FromIterator<A>: Sized {
    /// Creates a value from an iterator.
    ///
    /// See the [module-level documentation] for more.
    ///
    /// [module-level documentation]: index.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// use std::iter::FromIterator;
    ///
    /// let five_fives = std::iter::repeat(5).take(5);
    ///
    /// let v = Vec::from_iter(five_fives);
    ///
    /// assert_eq!(v, vec![5, 5, 5, 5, 5]);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn from_iter<T: IntoIterator<Item=A>>(iter: T) -> Self;
}

/// Conversion into an `Iterator`.
///
/// By implementing `IntoIterator` for a type, you define how it will be
/// converted to an iterator. This is common for types which describe a
/// collection of some kind.
///
/// One benefit of implementing `IntoIterator` is that your type will [work
/// with Rust's `for` loop syntax](index.html#for-loops-and-intoiterator).
///
/// See also: [`FromIterator`].
///
/// [`FromIterator`]: trait.FromIterator.html
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// let v = vec![1, 2, 3];
/// let mut iter = v.into_iter();
///
/// assert_eq!(Some(1), iter.next());
/// assert_eq!(Some(2), iter.next());
/// assert_eq!(Some(3), iter.next());
/// assert_eq!(None, iter.next());
/// ```
/// Implementing `IntoIterator` for your type:
///
/// ```
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // and we'll implement IntoIterator
/// impl IntoIterator for MyCollection {
///     type Item = i32;
///     type IntoIter = ::std::vec::IntoIter<Self::Item>;
///
///     fn into_iter(self) -> Self::IntoIter {
///         self.0.into_iter()
///     }
/// }
///
/// // Now we can make a new collection...
/// let mut c = MyCollection::new();
///
/// // ... add some stuff to it ...
/// c.add(0);
/// c.add(1);
/// c.add(2);
///
/// // ... and then turn it into an Iterator:
/// for (i, n) in c.into_iter().enumerate() {
///     assert_eq!(i as i32, n);
/// }
/// ```
///
/// It is common to use `IntoIterator` as a trait bound. This allows
/// the input collection type to change, so long as it is still an
/// iterator. Additional bounds can be specified by restricting on
/// `Item`:
///
/// ```rust
/// fn collect_as_strings<T>(collection: T) -> Vec<String>
///     where T: IntoIterator,
///           T::Item : std::fmt::Debug,
/// {
///     collection
///         .into_iter()
///         .map(|item| format!("{:?}", item))
///         .collect()
/// }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub trait IntoIterator {
    /// The type of the elements being iterated over.
    #[stable(feature = "rust1", since = "1.0.0")]
    type Item;

    /// Which kind of iterator are we turning this into?
    #[stable(feature = "rust1", since = "1.0.0")]
    type IntoIter: Iterator<Item=Self::Item>;

    /// Creates an iterator from a value.
    ///
    /// See the [module-level documentation] for more.
    ///
    /// [module-level documentation]: index.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let v = vec![1, 2, 3];
    /// let mut iter = v.into_iter();
    ///
    /// assert_eq!(Some(1), iter.next());
    /// assert_eq!(Some(2), iter.next());
    /// assert_eq!(Some(3), iter.next());
    /// assert_eq!(None, iter.next());
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn into_iter(self) -> Self::IntoIter;
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<I: Iterator> IntoIterator for I {
    type Item = I::Item;
    type IntoIter = I;

    fn into_iter(self) -> I {
        self
    }
}

/// Extend a collection with the contents of an iterator.
///
/// Iterators produce a series of values, and collections can also be thought
/// of as a series of values. The `Extend` trait bridges this gap, allowing you
/// to extend a collection by including the contents of that iterator. When
/// extending a collection with an already existing key, that entry is updated
/// or, in the case of collections that permit multiple entries with equal
/// keys, that entry is inserted.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// // You can extend a String with some chars:
/// let mut message = String::from("The first three letters are: ");
///
/// message.extend(&['a', 'b', 'c']);
///
/// assert_eq!("abc", &message[29..32]);
/// ```
///
/// Implementing `Extend`:
///
/// ```
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // since MyCollection has a list of i32s, we implement Extend for i32
/// impl Extend<i32> for MyCollection {
///
///     // This is a bit simpler with the concrete type signature: we can call
///     // extend on anything which can be turned into an Iterator which gives
///     // us i32s. Because we need i32s to put into MyCollection.
///     fn extend<T: IntoIterator<Item=i32>>(&mut self, iter: T) {
///
///         // The implementation is very straightforward: loop through the
///         // iterator, and add() each element to ourselves.
///         for elem in iter {
///             self.add(elem);
///         }
///     }
/// }
///
/// let mut c = MyCollection::new();
///
/// c.add(5);
/// c.add(6);
/// c.add(7);
///
/// // let's extend our collection with three more numbers
/// c.extend(vec![1, 2, 3]);
///
/// // we've added these elements onto the end
/// assert_eq!("MyCollection([5, 6, 7, 1, 2, 3])", format!("{:?}", c));
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub trait Extend<A> {
    /// Extends a collection with the contents of an iterator.
    ///
    /// As this is the only method for this trait, the [trait-level] docs
    /// contain more details.
    ///
    /// [trait-level]: trait.Extend.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// // You can extend a String with some chars:
    /// let mut message = String::from("abc");
    ///
    /// message.extend(['d', 'e', 'f'].iter());
    ///
    /// assert_eq!("abcdef", &message);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn extend<T: IntoIterator<Item=A>>(&mut self, iter: T);
}

#[stable(feature = "extend_for_unit", since = "1.28.0")]
impl Extend<()> for () {
    fn extend<T: IntoIterator<Item = ()>>(&mut self, iter: T) {
        iter.into_iter().for_each(drop)
    }
}
<commit_msg>doc(libcore) Fix CS<commit_after>/// Conversion from an `Iterator`.
///
/// By implementing `FromIterator` for a type, you define how it will be
/// created from an iterator. This is common for types which describe a
/// collection of some kind.
///
/// `FromIterator`'s [`from_iter`] is rarely called explicitly, and is instead
/// used through [`Iterator`]'s [`collect`] method. See [`collect`]'s
/// documentation for more examples.
///
/// [`from_iter`]: #tymethod.from_iter
/// [`Iterator`]: trait.Iterator.html
/// [`collect`]: trait.Iterator.html#method.collect
///
/// See also: [`IntoIterator`].
///
/// [`IntoIterator`]: trait.IntoIterator.html
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use std::iter::FromIterator;
///
/// let five_fives = std::iter::repeat(5).take(5);
///
/// let v = Vec::from_iter(five_fives);
///
/// assert_eq!(v, vec![5, 5, 5, 5, 5]);
/// ```
///
/// Using [`collect`] to implicitly use `FromIterator`:
///
/// ```
/// let five_fives = std::iter::repeat(5).take(5);
///
/// let v: Vec<i32> = five_fives.collect();
///
/// assert_eq!(v, vec![5, 5, 5, 5, 5]);
/// ```
///
/// Implementing `FromIterator` for your type:
///
/// ```
/// use std::iter::FromIterator;
///
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // and we'll implement FromIterator
/// impl FromIterator<i32> for MyCollection {
///     fn from_iter<I: IntoIterator<Item=i32>>(iter: I) -> Self {
///         let mut c = MyCollection::new();
///
///         for i in iter {
///             c.add(i);
///         }
///
///         c
///     }
/// }
///
/// // Now we can make a new iterator...
/// let iter = (0..5).into_iter();
///
/// // ... and make a MyCollection out of it
/// let c = MyCollection::from_iter(iter);
///
/// assert_eq!(c.0, vec![0, 1, 2, 3, 4]);
///
/// // collect works too!
///
/// let iter = (0..5).into_iter();
/// let c: MyCollection = iter.collect();
///
/// assert_eq!(c.0, vec![0, 1, 2, 3, 4]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
#[rustc_on_unimplemented(
    message="a collection of type `{Self}` cannot be built from an iterator \
             over elements of type `{A}`",
    label="a collection of type `{Self}` cannot be built from `std::iter::Iterator<Item={A}>`",
)]
pub trait FromIterator<A>: Sized {
    /// Creates a value from an iterator.
    ///
    /// See the [module-level documentation] for more.
    ///
    /// [module-level documentation]: index.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// use std::iter::FromIterator;
    ///
    /// let five_fives = std::iter::repeat(5).take(5);
    ///
    /// let v = Vec::from_iter(five_fives);
    ///
    /// assert_eq!(v, vec![5, 5, 5, 5, 5]);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn from_iter<T: IntoIterator<Item=A>>(iter: T) -> Self;
}

/// Conversion into an `Iterator`.
///
/// By implementing `IntoIterator` for a type, you define how it will be
/// converted to an iterator. This is common for types which describe a
/// collection of some kind.
///
/// One benefit of implementing `IntoIterator` is that your type will [work
/// with Rust's `for` loop syntax](index.html#for-loops-and-intoiterator).
///
/// See also: [`FromIterator`].
///
/// [`FromIterator`]: trait.FromIterator.html
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// let v = vec![1, 2, 3];
/// let mut iter = v.into_iter();
///
/// assert_eq!(Some(1), iter.next());
/// assert_eq!(Some(2), iter.next());
/// assert_eq!(Some(3), iter.next());
/// assert_eq!(None, iter.next());
/// ```
/// Implementing `IntoIterator` for your type:
///
/// ```
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // and we'll implement IntoIterator
/// impl IntoIterator for MyCollection {
///     type Item = i32;
///     type IntoIter = ::std::vec::IntoIter<Self::Item>;
///
///     fn into_iter(self) -> Self::IntoIter {
///         self.0.into_iter()
///     }
/// }
///
/// // Now we can make a new collection...
/// let mut c = MyCollection::new();
///
/// // ... add some stuff to it ...
/// c.add(0);
/// c.add(1);
/// c.add(2);
///
/// // ... and then turn it into an Iterator:
/// for (i, n) in c.into_iter().enumerate() {
///     assert_eq!(i as i32, n);
/// }
/// ```
///
/// It is common to use `IntoIterator` as a trait bound. This allows
/// the input collection type to change, so long as it is still an
/// iterator. Additional bounds can be specified by restricting on
/// `Item`:
///
/// ```rust
/// fn collect_as_strings<T>(collection: T) -> Vec<String>
///     where T: IntoIterator,
///           T::Item: std::fmt::Debug,
/// {
///     collection
///         .into_iter()
///         .map(|item| format!("{:?}", item))
///         .collect()
/// }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub trait IntoIterator {
    /// The type of the elements being iterated over.
    #[stable(feature = "rust1", since = "1.0.0")]
    type Item;

    /// Which kind of iterator are we turning this into?
    #[stable(feature = "rust1", since = "1.0.0")]
    type IntoIter: Iterator<Item=Self::Item>;

    /// Creates an iterator from a value.
    ///
    /// See the [module-level documentation] for more.
    ///
    /// [module-level documentation]: index.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// let v = vec![1, 2, 3];
    /// let mut iter = v.into_iter();
    ///
    /// assert_eq!(Some(1), iter.next());
    /// assert_eq!(Some(2), iter.next());
    /// assert_eq!(Some(3), iter.next());
    /// assert_eq!(None, iter.next());
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn into_iter(self) -> Self::IntoIter;
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<I: Iterator> IntoIterator for I {
    type Item = I::Item;
    type IntoIter = I;

    fn into_iter(self) -> I {
        self
    }
}

/// Extend a collection with the contents of an iterator.
///
/// Iterators produce a series of values, and collections can also be thought
/// of as a series of values. The `Extend` trait bridges this gap, allowing you
/// to extend a collection by including the contents of that iterator. When
/// extending a collection with an already existing key, that entry is updated
/// or, in the case of collections that permit multiple entries with equal
/// keys, that entry is inserted.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// // You can extend a String with some chars:
/// let mut message = String::from("The first three letters are: ");
///
/// message.extend(&['a', 'b', 'c']);
///
/// assert_eq!("abc", &message[29..32]);
/// ```
///
/// Implementing `Extend`:
///
/// ```
/// // A sample collection, that's just a wrapper over Vec<T>
/// #[derive(Debug)]
/// struct MyCollection(Vec<i32>);
///
/// // Let's give it some methods so we can create one and add things
/// // to it.
/// impl MyCollection {
///     fn new() -> MyCollection {
///         MyCollection(Vec::new())
///     }
///
///     fn add(&mut self, elem: i32) {
///         self.0.push(elem);
///     }
/// }
///
/// // since MyCollection has a list of i32s, we implement Extend for i32
/// impl Extend<i32> for MyCollection {
///
///     // This is a bit simpler with the concrete type signature: we can call
///     // extend on anything which can be turned into an Iterator which gives
///     // us i32s. Because we need i32s to put into MyCollection.
///     fn extend<T: IntoIterator<Item=i32>>(&mut self, iter: T) {
///
///         // The implementation is very straightforward: loop through the
///         // iterator, and add() each element to ourselves.
///         for elem in iter {
///             self.add(elem);
///         }
///     }
/// }
///
/// let mut c = MyCollection::new();
///
/// c.add(5);
/// c.add(6);
/// c.add(7);
///
/// // let's extend our collection with three more numbers
/// c.extend(vec![1, 2, 3]);
///
/// // we've added these elements onto the end
/// assert_eq!("MyCollection([5, 6, 7, 1, 2, 3])", format!("{:?}", c));
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub trait Extend<A> {
    /// Extends a collection with the contents of an iterator.
    ///
    /// As this is the only method for this trait, the [trait-level] docs
    /// contain more details.
    ///
    /// [trait-level]: trait.Extend.html
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// // You can extend a String with some chars:
    /// let mut message = String::from("abc");
    ///
    /// message.extend(['d', 'e', 'f'].iter());
    ///
    /// assert_eq!("abcdef", &message);
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    fn extend<T: IntoIterator<Item=A>>(&mut self, iter: T);
}

#[stable(feature = "extend_for_unit", since = "1.28.0")]
impl Extend<()> for () {
    fn extend<T: IntoIterator<Item = ()>>(&mut self, iter: T) {
        iter.into_iter().for_each(drop)
    }
}
<|endoftext|>


<commit_before>//! Solve systems of linear equations and invert matrices
//!
//! # Examples
//!
//! Solve `A * x = b`:
//!
//! ```
//! #[macro_use]
//! extern crate ndarray;
//! extern crate ndarray_linalg;
//!
//! use ndarray::prelude::*;
//! use ndarray_linalg::Solve;
//! # fn main() {
//!
//! let a: Array2<f64> = array![[3., 2., -1.], [2., -2., 4.], [-2., 1., -2.]];
//! let b: Array1<f64> = array![1., -2., 0.];
//! let x = a.solve_into(b).unwrap();
//! assert!(x.all_close(&array![1., -2., -2.], 1e-9));
//!
//! # }
//! ```
//!
//! There are also special functions for solving `A^T * x = b` and
//! `A^H * x = b`.
//!
//! If you are solving multiple systems of linear equations with the same
//! coefficient matrix `A`, it's faster to compute the LU factorization once at
//! the beginning than solving directly using `A`:
//!
//! ```
//! # extern crate ndarray;
//! # extern crate ndarray_linalg;
//! use ndarray::prelude::*;
//! use ndarray_linalg::*;
//! # fn main() {
//!
//! let a: Array2<f64> = random((3, 3));
//! let f = a.factorize_into().unwrap(); // LU factorize A (A is consumed)
//! for _ in 0..10 {
//!     let b: Array1<f64> = random(3);
//!     let x = f.solve_into(b).unwrap(); // Solve A * x = b using factorized L, U
//! }
//!
//! # }
//! ```

use ndarray::*;

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

pub use lapack_traits::{Pivot, Transpose};

/// An interface for solving systems of linear equations.
///
/// There are three groups of methods:
///
/// * `solve*` (normal) methods solve `A * x = b` for `x`.
/// * `solve_t*` (transpose) methods solve `A^T * x = b` for `x`.
/// * `solve_h*` (Hermitian conjugate) methods solve `A^H * x = b` for `x`.
///
/// Within each group, there are three methods that handle ownership differently:
///
/// * `*` methods take a reference to `b` and return `x` as a new array.
/// * `*_into` methods take ownership of `b`, store the result in it, and return it.
/// * `*_mut` methods take a mutable reference to `b` and store the result in that array.
///
/// If you plan to solve many equations with the same `A` matrix but different
/// `b` vectors, it's faster to factor the `A` matrix once using the
/// `Factorize` trait, and then solve using the `Factorized` struct.
pub trait Solve<A: Scalar> {
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;

    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_t_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_t_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;

    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_h_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_h_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
}

/// Represents the LU factorization of a matrix `A` as `A = P*L*U`.
pub struct Factorized<S: Data> {
    /// The factors `L` and `U`; the unit diagonal elements of `L` are not
    /// stored.
    pub a: ArrayBase<S, Ix2>,
    /// The pivot indices that define the permutation matrix `P`.
    pub ipiv: Pivot,
}

impl<A, S> Solve<A> for Factorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::No,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::Transpose,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::Hermite,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
}

impl<A, S> Solve<A> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_mut(rhs)
    }
    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_t_mut(rhs)
    }
    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_h_mut(rhs)
    }
}

/// An interface for computing LU factorizations of matrix refs.
pub trait Factorize<S: Data> {
    /// Computes the LU factorization `A = P*L*U`, where `P` is a permutation
    /// matrix.
    fn factorize(&self) -> Result<Factorized<S>>;
}

/// An interface for computing LU factorizations of matrices.
pub trait FactorizeInto<S: Data> {
    /// Computes the LU factorization `A = P*L*U`, where `P` is a permutation
    /// matrix.
    fn factorize_into(self) -> Result<Factorized<S>>;
}

impl<A, S> FactorizeInto<S> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    fn factorize_into(mut self) -> Result<Factorized<S>> {
        let ipiv = unsafe { A::lu(self.layout()?, self.as_allocated_mut()?)? };
        Ok(Factorized {
            a: self,
            ipiv: ipiv,
        })
    }
}

impl<A, Si> Factorize<OwnedRepr<A>> for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    fn factorize(&self) -> Result<Factorized<OwnedRepr<A>>> {
        let mut a: Array2<A> = replicate(self);
        let ipiv = unsafe { A::lu(a.layout()?, a.as_allocated_mut()?)? };
        Ok(Factorized { a: a, ipiv: ipiv })
    }
}

/// An interface for inverting matrix refs.
pub trait Inverse {
    type Output;
    /// Computes the inverse of the matrix.
    fn inv(&self) -> Result<Self::Output>;
}

/// An interface for inverting matrices.
pub trait InverseInto {
    type Output;
    /// Computes the inverse of the matrix.
    fn inv_into(self) -> Result<Self::Output>;
}

impl<A, S> InverseInto for Factorized<S>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    type Output = ArrayBase<S, Ix2>;

    fn inv_into(mut self) -> Result<ArrayBase<S, Ix2>> {
        unsafe {
            A::inv(
                self.a.square_layout()?,
                self.a.as_allocated_mut()?,
                &self.ipiv,
            )?
        };
        Ok(self.a)
    }
}

impl<A, S> Inverse for Factorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    type Output = Array2<A>;

    fn inv(&self) -> Result<Array2<A>> {
        let f = Factorized {
            a: replicate(&self.a),
            ipiv: self.ipiv.clone(),
        };
        f.inv_into()
    }
}

impl<A, S> InverseInto for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    type Output = Self;

    fn inv_into(self) -> Result<Self::Output> {
        let f = self.factorize_into()?;
        f.inv_into()
    }
}

impl<A, Si> Inverse for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    type Output = Array2<A>;

    fn inv(&self) -> Result<Self::Output> {
        let f = self.factorize()?;
        f.inv_into()
    }
}
<commit_msg>s/Factorized/LUFactorized/g<commit_after>//! Solve systems of linear equations and invert matrices
//!
//! # Examples
//!
//! Solve `A * x = b`:
//!
//! ```
//! #[macro_use]
//! extern crate ndarray;
//! extern crate ndarray_linalg;
//!
//! use ndarray::prelude::*;
//! use ndarray_linalg::Solve;
//! # fn main() {
//!
//! let a: Array2<f64> = array![[3., 2., -1.], [2., -2., 4.], [-2., 1., -2.]];
//! let b: Array1<f64> = array![1., -2., 0.];
//! let x = a.solve_into(b).unwrap();
//! assert!(x.all_close(&array![1., -2., -2.], 1e-9));
//!
//! # }
//! ```
//!
//! There are also special functions for solving `A^T * x = b` and
//! `A^H * x = b`.
//!
//! If you are solving multiple systems of linear equations with the same
//! coefficient matrix `A`, it's faster to compute the LU factorization once at
//! the beginning than solving directly using `A`:
//!
//! ```
//! # extern crate ndarray;
//! # extern crate ndarray_linalg;
//! use ndarray::prelude::*;
//! use ndarray_linalg::*;
//! # fn main() {
//!
//! let a: Array2<f64> = random((3, 3));
//! let f = a.factorize_into().unwrap(); // LU factorize A (A is consumed)
//! for _ in 0..10 {
//!     let b: Array1<f64> = random(3);
//!     let x = f.solve_into(b).unwrap(); // Solve A * x = b using factorized L, U
//! }
//!
//! # }
//! ```

use ndarray::*;

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

pub use lapack_traits::{Pivot, Transpose};

/// An interface for solving systems of linear equations.
///
/// There are three groups of methods:
///
/// * `solve*` (normal) methods solve `A * x = b` for `x`.
/// * `solve_t*` (transpose) methods solve `A^T * x = b` for `x`.
/// * `solve_h*` (Hermitian conjugate) methods solve `A^H * x = b` for `x`.
///
/// Within each group, there are three methods that handle ownership differently:
///
/// * `*` methods take a reference to `b` and return `x` as a new array.
/// * `*_into` methods take ownership of `b`, store the result in it, and return it.
/// * `*_mut` methods take a mutable reference to `b` and store the result in that array.
///
/// If you plan to solve many equations with the same `A` matrix but different
/// `b` vectors, it's faster to factor the `A` matrix once using the
/// `Factorize` trait, and then solve using the `LUFactorized` struct.
pub trait Solve<A: Scalar> {
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;

    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_t_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_t_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^T * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_t_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;

    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solve_h_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solve_h_mut(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A^H * x = b` where `A` is `self`, `b`
    /// is the argument, and `x` is the successful result.
    fn solve_h_mut<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
}

/// Represents the LU factorization of a matrix `A` as `A = P*L*U`.
pub struct LUFactorized<S: Data> {
    /// The factors `L` and `U`; the unit diagonal elements of `L` are not
    /// stored.
    pub a: ArrayBase<S, Ix2>,
    /// The pivot indices that define the permutation matrix `P`.
    pub ipiv: Pivot,
}

impl<A, S> Solve<A> for LUFactorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::No,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::Transpose,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solve(
                self.a.square_layout()?,
                Transpose::Hermite,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
}

impl<A, S> Solve<A> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solve_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_mut(rhs)
    }
    fn solve_t_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_t_mut(rhs)
    }
    fn solve_h_mut<'a, Sb>(&self, mut rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        let f = self.factorize()?;
        f.solve_h_mut(rhs)
    }
}

/// An interface for computing LU factorizations of matrix refs.
pub trait Factorize<S: Data> {
    /// Computes the LU factorization `A = P*L*U`, where `P` is a permutation
    /// matrix.
    fn factorize(&self) -> Result<LUFactorized<S>>;
}

/// An interface for computing LU factorizations of matrices.
pub trait FactorizeInto<S: Data> {
    /// Computes the LU factorization `A = P*L*U`, where `P` is a permutation
    /// matrix.
    fn factorize_into(self) -> Result<LUFactorized<S>>;
}

impl<A, S> FactorizeInto<S> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    fn factorize_into(mut self) -> Result<LUFactorized<S>> {
        let ipiv = unsafe { A::lu(self.layout()?, self.as_allocated_mut()?)? };
        Ok(LUFactorized {
            a: self,
            ipiv: ipiv,
        })
    }
}

impl<A, Si> Factorize<OwnedRepr<A>> for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    fn factorize(&self) -> Result<LUFactorized<OwnedRepr<A>>> {
        let mut a: Array2<A> = replicate(self);
        let ipiv = unsafe { A::lu(a.layout()?, a.as_allocated_mut()?)? };
        Ok(LUFactorized { a: a, ipiv: ipiv })
    }
}

/// An interface for inverting matrix refs.
pub trait Inverse {
    type Output;
    /// Computes the inverse of the matrix.
    fn inv(&self) -> Result<Self::Output>;
}

/// An interface for inverting matrices.
pub trait InverseInto {
    type Output;
    /// Computes the inverse of the matrix.
    fn inv_into(self) -> Result<Self::Output>;
}

impl<A, S> InverseInto for LUFactorized<S>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    type Output = ArrayBase<S, Ix2>;

    fn inv_into(mut self) -> Result<ArrayBase<S, Ix2>> {
        unsafe {
            A::inv(
                self.a.square_layout()?,
                self.a.as_allocated_mut()?,
                &self.ipiv,
            )?
        };
        Ok(self.a)
    }
}

impl<A, S> Inverse for LUFactorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    type Output = Array2<A>;

    fn inv(&self) -> Result<Array2<A>> {
        let f = LUFactorized {
            a: replicate(&self.a),
            ipiv: self.ipiv.clone(),
        };
        f.inv_into()
    }
}

impl<A, S> InverseInto for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    type Output = Self;

    fn inv_into(self) -> Result<Self::Output> {
        let f = self.factorize_into()?;
        f.inv_into()
    }
}

impl<A, Si> Inverse for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    type Output = Array2<A>;

    fn inv(&self) -> Result<Self::Output> {
        let f = self.factorize()?;
        f.inv_into()
    }
}
<|endoftext|>