<commit_before>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[feature(macro_rules)];

extern crate extra;
extern crate getopts;
extern crate sync;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use extra::time::Timespec;
use sync::{Arc, Future};

#[path = "../util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("", "help", "display this help and exit"),
        getopts::optflag("", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    } else if matches.opt_present("version") {
        println!("du version: {}", VERSION);
        return
    }

    let options = Options{
        all: matches.opt_present("all"),
        max_depth: match (matches.opt_present("summarize"), matches.opt_str("max-depth")) {
            (true, Some(s)) => match from_str::<uint>(s) {
                Some(_) => {
                    println!("du: warning: summarizing conflicts with --max-depth={:s}", s);
                    return
                },
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (true, None) => Some(0),
            (false, Some(s)) => match from_str::<uint>(s) {
                Some(u) => Some(u),
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (false, None) => None
        },
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = matches.free.clone();
    let strs = match strs.is_empty() {
        true => ~[~"./"],
        false => strs
    };

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    extra::time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<commit_msg>du: add -V and -h<commit_after>#[crate_id(name="du", vers="1.0.0", author="Derek Chiang")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Derek Chiang <derekchiang93@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

#[feature(macro_rules)];

extern crate extra;
extern crate getopts;
extern crate sync;

use std::os;
use std::io::fs;
use std::io::FileStat;
use std::option::Option;
use std::path::Path;
use extra::time::Timespec;
use sync::{Arc, Future};

#[path = "../util.rs"]
mod util;

static NAME: &'static str = "du";
static VERSION: &'static str = "1.0.0";

struct Options {
    all: bool,
    max_depth: Option<uint>,
    total: bool,
    separate_dirs: bool,
}

fn du(path: &Path, options_arc: Arc<Options>, depth: uint) -> ~[Arc<FileStat>] {
    let mut stats = ~[];
    let mut futures = ~[];
    let options = options_arc.get();
    let mut my_stat = safe_unwrap!(path.stat());

    for f in safe_unwrap!(fs::readdir(path)).move_iter() {
        match f.is_file() {
            true => {
                let stat = safe_unwrap!(f.stat());
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
                if options.all {
                    stats.push(Arc::new(stat))
                }
            }
            false => {
                let oa_clone = options_arc.clone();
                futures.push(Future::spawn(proc() { du(&f, oa_clone, depth + 1) }))
            }
        }
    }

    for future in futures.mut_iter() {
        for stat_arc in future.get().move_rev_iter() {
            let stat = stat_arc.get();
            if !options.separate_dirs && stat.path.dir_path() == my_stat.path {
                my_stat.size += stat.size;
                my_stat.unstable.blocks += stat.unstable.blocks;
            }
            if options.max_depth == None || depth < options.max_depth.unwrap() {
                stats.push(stat_arc.clone());
            }
        }
    }

    stats.push(Arc::new(my_stat));

    return stats;
}

fn main() {
    let args = os::args();
    let program = args[0].clone();
    let opts = ~[
        // In task
        getopts::optflag("a", "all", " write counts for all files, not just directories"),
        // In main
        getopts::optflag("", "apparent-size", "print apparent sizes,  rather  than  disk  usage;
            although  the apparent  size is usually smaller, it may be larger due to holes
            in ('sparse') files, internal  fragmentation,  indirect  blocks, and the like"),
        // In main
        getopts::optopt("B", "block-size", "scale sizes  by  SIZE before printing them.
            E.g., '-BM' prints sizes in units of 1,048,576 bytes.  See SIZE format below.",
            "SIZE"),
        // In main
        getopts::optflag("b", "bytes", "equivalent to '--apparent-size --block-size=1'"),
        // In main
        getopts::optflag("c", "total", "produce a grand total"),
        // In task
        // getopts::optflag("D", "dereference-args", "dereference only symlinks that are listed
        //     on the command line"),
        // In main
        // getopts::optopt("", "files0-from", "summarize disk usage of the NUL-terminated file
        //                   names specified in file F;
        //                   If F is - then read names from standard input", "F"),
        // // In task
        // getopts::optflag("H", "", "equivalent to --dereference-args (-D)"),
        // In main
        getopts::optflag("h", "human-readable", "print sizes in human readable format (e.g., 1K 234M 2G)"),
        // In main
        getopts::optflag("", "si", "like -h, but use powers of 1000 not 1024"),
        // In main
        getopts::optflag("k", "", "like --block-size=1K"),
        // In task
        getopts::optflag("l", "count-links", "count sizes many times if hard linked"),
        // // In main
        getopts::optflag("m", "", "like --block-size=1M"),
        // // In task
        // getopts::optflag("L", "dereference", "dereference all symbolic links"),
        // // In task
        // getopts::optflag("P", "no-dereference", "don't follow any symbolic links (this is the default)"),
        // // In main
        getopts::optflag("0", "null", "end each output line with 0 byte rather than newline"),
        // In main
        getopts::optflag("S", "separate-dirs", "do not include size of subdirectories"),
        // In main
        getopts::optflag("s", "summarize", "display only a total for each argument"),
        // // In task
        // getopts::optflag("x", "one-file-system", "skip directories on different file systems"),
        // // In task
        // getopts::optopt("X", "exclude-from", "exclude files that match any pattern in FILE", "FILE"),
        // // In task
        // getopts::optopt("", "exclude", "exclude files that match PATTERN", "PATTERN"),
        // In main
        getopts::optopt("d", "max-depth", "print the total for a directory (or file, with --all)
            only if it is N or fewer levels below the command
            line argument;  --max-depth=0 is the same as --summarize", "N"),
        // In main
        getopts::optflagopt("", "time", "show time of the last modification of any file in the
            directory, or any of its subdirectories.  If WORD is given, show time as WORD instead of modification time:
            atime, access, use, ctime or status", "WORD"),
        // In main
        getopts::optopt("", "time-style", "show times using style STYLE:
            full-iso, long-iso, iso, +FORMAT FORMAT is interpreted like 'date'", "STYLE"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(f) => {
            show_error!(1, "Invalid options\n{}", f.to_err_msg());
            return
        }
    };

    if matches.opt_present("help") {
        println!("du {} - estimate file space usage", VERSION);
        println!("");
        println!("Usage:");
        println!("  {0:s} [OPTION]... [FILE]...", program);
        println!("  {0:s} [OPTION]... --files0-from=F", program);
        println!("");
        println!("{}", getopts::usage("Summarize disk usage of each FILE, recursively for directories.", opts));
        println!("Display  values  are  in  units  of  the  first  available  SIZE from
--block-size,  and the DU_BLOCK_SIZE, BLOCK_SIZE and BLOCKSIZE environ‐
ment variables.  Otherwise, units default to  1024  bytes  (or  512  if
POSIXLY_CORRECT is set).

SIZE  is  an  integer and optional unit (example: 10M is 10*1024*1024).
Units are K, M, G, T, P, E, Z, Y (powers of 1024) or KB, MB, ...  (pow‐
ers of 1000).");
        return
    } else if matches.opt_present("version") {
        println!("du version: {}", VERSION);
        return
    }

    let options = Options{
        all: matches.opt_present("all"),
        max_depth: match (matches.opt_present("summarize"), matches.opt_str("max-depth")) {
            (true, Some(s)) => match from_str::<uint>(s) {
                Some(_) => {
                    println!("du: warning: summarizing conflicts with --max-depth={:s}", s);
                    return
                },
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (true, None) => Some(0),
            (false, Some(s)) => match from_str::<uint>(s) {
                Some(u) => Some(u),
                None => {
                    println!("du: invalid maximum depth '{:s}'", s);
                    return
                }
            },
            (false, None) => None
        },
        total: matches.opt_present("total"),
        separate_dirs: matches.opt_present("S"),
    };

    let strs = matches.free.clone();
    let strs = match strs.is_empty() {
        true => ~[~"./"],
        false => strs
    };

    let options_arc = Arc::new(options);

    let MB = match matches.opt_present("si") {
        true => 1000 * 1000,
        false => 1024 * 1024,
    };
    let KB = match matches.opt_present("si") {
        true => 1000,
        false => 1024,
    };

    let block_size = match matches.opt_str("block-size") {
        Some(s) => {
            let mut found_number = false;
            let mut found_letter = false;
            let mut numbers = ~[];
            let mut letters = ~[];
            for c in s.chars() {
                if found_letter && c.is_digit() || !found_number && !c.is_digit() {
                    println!("du: invalid --block-size argument '{}'", s);
                    return
                } else if c.is_digit() {
                    found_number = true;
                    numbers.push(c as u8);
                } else if c.is_alphabetic() {
                    found_letter = true;
                    letters.push(c);
                }
            }
            let number = std::uint::parse_bytes(numbers, 10).unwrap();
            let multiple = match std::str::from_chars(letters).as_slice() {
                "K" => 1024, "M" => 1024 * 1024, "G" => 1024 * 1024 * 1024,
                "T" => 1024 * 1024 * 1024 * 1024, "P" => 1024 * 1024 * 1024 * 1024 * 1024,
                "E" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Z" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "Y" => 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024 * 1024,
                "KB" => 1000, "MB" => 1000 * 1000, "GB" => 1000 * 1000 * 1000,
                "TB" => 1000 * 1000 * 1000 * 1000, "PB" => 1000 * 1000 * 1000 * 1000 * 1000,
                "EB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "ZB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                "YB" => 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000,
                _ => {
                    println!("du: invalid --block-size argument '{}'", s); return
                }
            };
            number * multiple
        },
        None => 1024
    };

    let convert_size = |size: u64| -> ~str {
        if matches.opt_present("human-readable") || matches.opt_present("si") {
            if size > MB {
                format!("{:.1f}M", (size as f64) / (MB as f64))
            } else if size > KB {
                format!("{:.1f}K", (size as f64) / (KB as f64))
            } else {
                format!("{}B", size)
            }
        } else if matches.opt_present("k") {
            format!("{}", ((size as f64) / (KB as f64)).ceil())
        } else if matches.opt_present("m") {
            format!("{}", ((size as f64) / (MB as f64)).ceil())
        } else {
            format!("{}", ((size as f64) / (block_size as f64)).ceil())
        }
    };

    let time_format_str = match matches.opt_str("time-style") {
        Some(s) => {
            match s.as_slice() {
                "full-iso" => "%Y-%m-%d %H:%M:%S.%f %z",
                "long-iso" => "%Y-%m-%d %H:%M",
                "iso" => "%Y-%m-%d",
                _ => {
                    println!("
du: invalid argument 'awdwa' for 'time style'
Valid arguments are:
- 'full-iso'
- 'long-iso'
- 'iso'
Try 'du --help' for more information.");
                    return
                }
            }
        },
        None => "%Y-%m-%d %H:%M"
    };

    let line_separator = match matches.opt_present("0") {
        true => "\0",
        false => "\n",
    };

    let mut grand_total = 0;
    for path_str in strs.iter() {
        let path = Path::new(path_str.clone());
        let iter = du(&path, options_arc.clone(), 0).move_iter();
        let (_, len) = iter.size_hint();
        let len = len.unwrap();
        for (index, stat_arc) in iter.enumerate() {
            let stat = stat_arc.get();
            let size = match matches.opt_present("apparent-size") {
                true => stat.unstable.nlink * stat.size,
                // C's stat is such that each block is assume to be 512 bytes
                // See: http://linux.die.net/man/2/stat
                false => stat.unstable.blocks * 512,
            };
            if matches.opt_present("time") {
                let time_str = {
                    let (secs, nsecs) = {
                        let time = match matches.opt_str("time") {
                            Some(s) => match s.as_slice() {
                                "accessed" => stat.accessed,
                                "created" => stat.created,
                                "modified" => stat.modified,
                                _ => {
                                    println!("du: invalid argument 'modified' for '--time'
    Valid arguments are:
      - 'accessed', 'created', 'modified'
    Try 'du --help' for more information.");
                                    return
                                }
                            },
                            None => stat.modified
                        };
                        ((time / 1000) as i64, (time % 1000 * 1000000) as i32)
                    };
                    let time_spec = Timespec::new(secs, nsecs);
                    extra::time::at(time_spec).strftime(time_format_str)
                };
                print!("{:<10} {:<30} {}", convert_size(size), time_str, stat.path.display());
            } else {
                print!("{:<10} {}", convert_size(size), stat.path.display());
            }
            print!("{}", line_separator);
            if options.total && index == (len - 1) {
                // The last element will be the total size of the the path under
                // path_str.  We add it to the grand total.
                grand_total += size;
            }
        }
    }

    if options.total {
        print!("{:<10} total", convert_size(grand_total));
        print!("{}", line_separator);
    }
}
<|endoftext|>


<commit_before>use std::io::Cursor;

use crate::crypto::{PublicKey, ScriptHash, KeyPair, SigPair};
use crate::script::Script;
use crate::serializer::*;
use crate::asset::Asset;

#[macro_use] mod util;

pub mod tx_type;
pub use self::tx_type::*;

pub trait EncodeTx {
    fn encode(&self, v: &mut Vec<u8>);
}

pub trait DecodeTx<T> {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<T>;
}

pub trait SignTx {
    fn sign(&self, key_pair: &KeyPair) -> SigPair;
    fn append_sign(&mut self, key_pair: &KeyPair);
}

#[derive(Debug, Clone)]
pub enum TxVariant {
    RewardTx(RewardTx),
    BondTx(BondTx),
    TransferTx(TransferTx)
}

impl TxVariant {
    pub fn encode(&self, v: &mut Vec<u8>) {
        match self {
            TxVariant::RewardTx(tx) => { tx.encode(v) },
            TxVariant::BondTx(tx) => { tx.encode(v) },
            TxVariant::TransferTx(tx) => { tx.encode(v) }
        };
    }

    pub fn encode_with_sigs(&self, v: &mut Vec<u8>) {
        macro_rules! encode_sigs {
            ($name:expr, $vec:expr) => {
                {
                    $vec.push_u16($name.signature_pairs.len() as u16);
                    for sig in &$name.signature_pairs { $vec.push_sig_pair(sig) };
                    $name.encode($vec);
                }
            }
        }

        match self {
            TxVariant::RewardTx(tx) => { encode_sigs!(tx, v) },
            TxVariant::BondTx(tx) => { encode_sigs!(tx, v) },
            TxVariant::TransferTx(tx) => { encode_sigs!(tx, v) }
        };
    }

    pub fn decode_with_sigs(cur: &mut Cursor<&[u8]>) -> Option<TxVariant> {
        let sigs = {
            let len = cur.take_u16().ok()?;
            let mut vec = Vec::with_capacity(len as usize);
            for _ in 0..len { vec.push(cur.take_sig_pair().ok()?) };
            vec
        };
        let mut base = Tx::decode_base(cur)?;
        base.signature_pairs = sigs;
        match base.tx_type {
            TxType::REWARD => Some(TxVariant::RewardTx(RewardTx::decode(cur, base)?)),
            TxType::BOND => Some(TxVariant::BondTx(BondTx::decode(cur, base)?)),
            TxType::TRANSFER => Some(TxVariant::TransferTx(TransferTx::decode(cur, base)?))
        }
    }
}

impl ::std::ops::Deref for TxVariant {
    type Target = Tx;

    fn deref(&self) -> &Self::Target {
        match self {
            TxVariant::RewardTx(tx) => &tx.base,
            TxVariant::BondTx(tx) => &tx.base,
            TxVariant::TransferTx(tx) => &tx.base
        }
    }
}

#[derive(Debug, Clone)]
pub struct Tx {
    pub tx_type: TxType,
    pub timestamp: u32,
    pub fee: Asset,
    pub signature_pairs: Vec<SigPair>
}

impl Tx {
    fn encode_base(&self, v: &mut Vec<u8>) {
        v.push(self.tx_type as u8);
        v.push_u32(self.timestamp);
        v.push_asset(&self.fee);
    }

    fn decode_base(cur: &mut Cursor<&[u8]>) -> Option<Tx> {
        let tx_type = match cur.take_u8().ok()? {
            t if t == TxType::REWARD as u8 => TxType::REWARD,
            t if t == TxType::BOND as u8 => TxType::BOND,
            t if t == TxType::TRANSFER as u8 => TxType::TRANSFER,
            _ => return None
        };
        let timestamp = cur.take_u32().ok()?;
        let fee = cur.take_asset().ok()?;

        Some(Tx {
            tx_type,
            timestamp,
            fee,
            signature_pairs: Vec::new()
        })
    }
}

#[derive(Debug, Clone)]
pub struct RewardTx {
    pub base: Tx,
    pub to: ScriptHash,
    pub rewards: Vec<Asset>
}

impl EncodeTx for RewardTx {
    fn encode(&self, v: &mut Vec<u8>) {
        debug_assert_eq!(self.base.signature_pairs.len(), 0);
        self.encode_base(v);
        v.push_script_hash(&self.to);
        v.push_u32(self.rewards.len() as u32);
        for r in &self.rewards { v.push_asset(r) };
    }
}

impl DecodeTx<RewardTx> for RewardTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<RewardTx> {
        assert_eq!(tx.tx_type, TxType::REWARD);
        let key = cur.take_script_hash().ok()?;

        let len = cur.take_u32().ok()?;
        let mut rewards = Vec::with_capacity(len as usize);
        for _ in 0..len {
            rewards.push(cur.take_asset().ok()?);
        }

        Some(RewardTx {
            base: tx,
            to: key,
            rewards
        })
    }
}

#[derive(Debug, Clone)]
pub struct BondTx {
    pub base: Tx,
    pub minter: PublicKey, // Key that signs blocks
    pub staker: ScriptHash, // Hot wallet that receives rewards and stakes its balance
    pub stake_amt: Asset,
    pub bond_fee: Asset
}

impl EncodeTx for BondTx {
    fn encode(&self, v: &mut Vec<u8>) {
        self.encode_base(v);
        v.push_pub_key(&self.minter);
        v.push_script_hash(&self.staker);
        v.push_asset(&self.stake_amt);
        v.push_asset(&self.bond_fee);
    }
}

impl DecodeTx<BondTx> for BondTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<BondTx> {
        assert_eq!(tx.tx_type, TxType::BOND);
        let minter = cur.take_pub_key().ok()?;
        let staker = cur.take_script_hash().ok()?;
        let stake_amt = cur.take_asset().ok()?;
        let bond_fee = cur.take_asset().ok()?;
        Some(BondTx {
            base: tx,
            minter,
            staker,
            stake_amt,
            bond_fee
        })
    }
}

#[derive(Debug, Clone)]
pub struct TransferTx {
    pub base: Tx,
    pub from: ScriptHash,
    pub to: ScriptHash,
    pub script: Script,
    pub amount: Asset,
    pub memo: Vec<u8>
}

impl EncodeTx for TransferTx {
    fn encode(&self, v: &mut Vec<u8>) {
        self.encode_base(v);
        v.push_script_hash(&self.from);
        v.push_script_hash(&self.to);
        v.push_bytes(&self.script);
        v.push_asset(&self.amount);
        v.push_bytes(&self.memo);
    }
}

impl DecodeTx<TransferTx> for TransferTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<TransferTx> {
        assert_eq!(tx.tx_type, TxType::TRANSFER);
        let from = cur.take_script_hash().ok()?;
        let to = cur.take_script_hash().ok()?;
        let script = Script::new(cur.take_bytes().ok()?);
        let amount = cur.take_asset().ok()?;
        let memo = cur.take_bytes().ok()?;
        Some(TransferTx {
            base: tx,
            from,
            to,
            script,
            amount,
            memo
        })
    }
}

tx_deref!(RewardTx);
tx_deref!(BondTx);
tx_deref!(TransferTx);

tx_sign!(BondTx);
tx_sign!(TransferTx);

#[cfg(test)]
mod tests {
    use ::std::str::FromStr;

    use crate::crypto;
    use super::*;

    macro_rules! cmp_base_tx {
        ($id:ident, $ty:expr, $ts:expr, $fee:expr) => {
            assert_eq!($id.tx_type, $ty);
            assert_eq!($id.timestamp, $ts);
            assert_eq!($id.fee.to_string(), $fee);
        }
    }

    #[test]
    fn test_encode_tx_with_sigs() {
        let to = crypto::KeyPair::gen_keypair();
        let reward_tx = TxVariant::RewardTx(RewardTx {
            base: Tx {
                tx_type: TxType::REWARD,
                timestamp: 123,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            to: to.0.into(),
            rewards: vec![get_asset("1.50 GOLD"), get_asset("1.0 SILVER")]
        });

        let mut v = vec![];
        reward_tx.encode_with_sigs(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        TxVariant::decode_with_sigs(&mut c).unwrap();
    }

    #[test]
    fn test_encode_reward() {
        let to = crypto::KeyPair::gen_keypair();
        let reward_tx = RewardTx {
            base: Tx {
                tx_type: TxType::REWARD,
                timestamp: 123,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            to: to.0.into(),
            rewards: vec![get_asset("1.50 GOLD"), get_asset("1.0 SILVER")]
        };

        let mut v = vec![];
        reward_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = RewardTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::REWARD, 123, "123 GOLD");
        assert_eq!(reward_tx.to, dec.to);
        assert_eq!(reward_tx.rewards.len(), dec.rewards.len());
        assert_eq!(reward_tx.rewards[0].to_string(), dec.rewards[0].to_string());
        assert_eq!(reward_tx.rewards[1].to_string(), dec.rewards[1].to_string());
    }

    #[test]
    fn test_encode_bond() {
        let minter = crypto::KeyPair::gen_keypair();
        let staker = crypto::KeyPair::gen_keypair();
        let bond_tx = BondTx {
            base: Tx {
                tx_type: TxType::BOND,
                timestamp: 1230,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            minter: minter.0,
            staker: staker.0.into(),
            stake_amt: get_asset("1.0456 GOLD"),
            bond_fee: get_asset("1.00000000 GOLD")
        };

        let mut v = vec![];
        bond_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = BondTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::BOND, 1230, "123 GOLD");
        assert_eq!(bond_tx.minter, dec.minter);
        assert_eq!(bond_tx.staker, dec.staker);
        assert_eq!(bond_tx.stake_amt.to_string(), dec.stake_amt.to_string());
        assert_eq!(bond_tx.bond_fee.to_string(), dec.bond_fee.to_string());
    }

    #[test]
    fn test_encode_transfer() {
        let from = crypto::KeyPair::gen_keypair();
        let to = crypto::KeyPair::gen_keypair();
        let transfer_tx = TransferTx {
            base: Tx {
                tx_type: TxType::TRANSFER,
                timestamp: 1234567890,
                fee: get_asset("1.23 GOLD"),
                signature_pairs: vec![]
            },
            from: from.0.into(),
            to: to.0.into(),
            script: Script::new(vec![1, 2, 3, 4]),
            amount: get_asset("1.0456 GOLD"),
            memo: Vec::from(String::from("Hello world!").as_bytes())
        };

        let mut v = vec![];
        transfer_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = TransferTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::TRANSFER, 1234567890, "1.23 GOLD");
        assert_eq!(transfer_tx.from, dec.from);
        assert_eq!(transfer_tx.to, dec.to);
        assert_eq!(transfer_tx.script, Script::new(vec![1, 2, 3, 4]));
        assert_eq!(transfer_tx.amount.to_string(), dec.amount.to_string());
        assert_eq!(transfer_tx.memo, dec.memo);
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}
<commit_msg>Implement Into for CoW TxVariants<commit_after>use std::borrow::Cow;
use std::io::Cursor;

use crate::crypto::{PublicKey, ScriptHash, KeyPair, SigPair};
use crate::script::Script;
use crate::serializer::*;
use crate::asset::Asset;

#[macro_use] mod util;

pub mod tx_type;
pub use self::tx_type::*;

pub trait EncodeTx {
    fn encode(&self, v: &mut Vec<u8>);
}

pub trait DecodeTx<T> {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<T>;
}

pub trait SignTx {
    fn sign(&self, key_pair: &KeyPair) -> SigPair;
    fn append_sign(&mut self, key_pair: &KeyPair);
}

#[derive(Debug, Clone)]
pub enum TxVariant {
    RewardTx(RewardTx),
    BondTx(BondTx),
    TransferTx(TransferTx)
}

impl TxVariant {
    pub fn encode(&self, v: &mut Vec<u8>) {
        match self {
            TxVariant::RewardTx(tx) => { tx.encode(v) },
            TxVariant::BondTx(tx) => { tx.encode(v) },
            TxVariant::TransferTx(tx) => { tx.encode(v) }
        };
    }

    pub fn encode_with_sigs(&self, v: &mut Vec<u8>) {
        macro_rules! encode_sigs {
            ($name:expr, $vec:expr) => {
                {
                    $vec.push_u16($name.signature_pairs.len() as u16);
                    for sig in &$name.signature_pairs { $vec.push_sig_pair(sig) };
                    $name.encode($vec);
                }
            }
        }

        match self {
            TxVariant::RewardTx(tx) => { encode_sigs!(tx, v) },
            TxVariant::BondTx(tx) => { encode_sigs!(tx, v) },
            TxVariant::TransferTx(tx) => { encode_sigs!(tx, v) }
        };
    }

    pub fn decode_with_sigs(cur: &mut Cursor<&[u8]>) -> Option<TxVariant> {
        let sigs = {
            let len = cur.take_u16().ok()?;
            let mut vec = Vec::with_capacity(len as usize);
            for _ in 0..len { vec.push(cur.take_sig_pair().ok()?) };
            vec
        };
        let mut base = Tx::decode_base(cur)?;
        base.signature_pairs = sigs;
        match base.tx_type {
            TxType::REWARD => Some(TxVariant::RewardTx(RewardTx::decode(cur, base)?)),
            TxType::BOND => Some(TxVariant::BondTx(BondTx::decode(cur, base)?)),
            TxType::TRANSFER => Some(TxVariant::TransferTx(TransferTx::decode(cur, base)?))
        }
    }
}

impl ::std::ops::Deref for TxVariant {
    type Target = Tx;

    fn deref(&self) -> &Self::Target {
        match self {
            TxVariant::RewardTx(tx) => &tx.base,
            TxVariant::BondTx(tx) => &tx.base,
            TxVariant::TransferTx(tx) => &tx.base
        }
    }
}

impl<'a> Into<Cow<'a, TxVariant>> for TxVariant {
    fn into(self) -> Cow<'a, TxVariant> {
        Cow::Owned(self)
    }
}

impl<'a> Into<Cow<'a, TxVariant>> for &'a TxVariant {
    fn into(self) -> Cow<'a, TxVariant> {
        Cow::Borrowed(self)
    }
}

#[derive(Debug, Clone)]
pub struct Tx {
    pub tx_type: TxType,
    pub timestamp: u32,
    pub fee: Asset,
    pub signature_pairs: Vec<SigPair>
}

impl Tx {
    fn encode_base(&self, v: &mut Vec<u8>) {
        v.push(self.tx_type as u8);
        v.push_u32(self.timestamp);
        v.push_asset(&self.fee);
    }

    fn decode_base(cur: &mut Cursor<&[u8]>) -> Option<Tx> {
        let tx_type = match cur.take_u8().ok()? {
            t if t == TxType::REWARD as u8 => TxType::REWARD,
            t if t == TxType::BOND as u8 => TxType::BOND,
            t if t == TxType::TRANSFER as u8 => TxType::TRANSFER,
            _ => return None
        };
        let timestamp = cur.take_u32().ok()?;
        let fee = cur.take_asset().ok()?;

        Some(Tx {
            tx_type,
            timestamp,
            fee,
            signature_pairs: Vec::new()
        })
    }
}

#[derive(Debug, Clone)]
pub struct RewardTx {
    pub base: Tx,
    pub to: ScriptHash,
    pub rewards: Vec<Asset>
}

impl EncodeTx for RewardTx {
    fn encode(&self, v: &mut Vec<u8>) {
        debug_assert_eq!(self.base.signature_pairs.len(), 0);
        self.encode_base(v);
        v.push_script_hash(&self.to);
        v.push_u32(self.rewards.len() as u32);
        for r in &self.rewards { v.push_asset(r) };
    }
}

impl DecodeTx<RewardTx> for RewardTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<RewardTx> {
        assert_eq!(tx.tx_type, TxType::REWARD);
        let key = cur.take_script_hash().ok()?;

        let len = cur.take_u32().ok()?;
        let mut rewards = Vec::with_capacity(len as usize);
        for _ in 0..len {
            rewards.push(cur.take_asset().ok()?);
        }

        Some(RewardTx {
            base: tx,
            to: key,
            rewards
        })
    }
}

#[derive(Debug, Clone)]
pub struct BondTx {
    pub base: Tx,
    pub minter: PublicKey, // Key that signs blocks
    pub staker: ScriptHash, // Hot wallet that receives rewards and stakes its balance
    pub stake_amt: Asset,
    pub bond_fee: Asset
}

impl EncodeTx for BondTx {
    fn encode(&self, v: &mut Vec<u8>) {
        self.encode_base(v);
        v.push_pub_key(&self.minter);
        v.push_script_hash(&self.staker);
        v.push_asset(&self.stake_amt);
        v.push_asset(&self.bond_fee);
    }
}

impl DecodeTx<BondTx> for BondTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<BondTx> {
        assert_eq!(tx.tx_type, TxType::BOND);
        let minter = cur.take_pub_key().ok()?;
        let staker = cur.take_script_hash().ok()?;
        let stake_amt = cur.take_asset().ok()?;
        let bond_fee = cur.take_asset().ok()?;
        Some(BondTx {
            base: tx,
            minter,
            staker,
            stake_amt,
            bond_fee
        })
    }
}

#[derive(Debug, Clone)]
pub struct TransferTx {
    pub base: Tx,
    pub from: ScriptHash,
    pub to: ScriptHash,
    pub script: Script,
    pub amount: Asset,
    pub memo: Vec<u8>
}

impl EncodeTx for TransferTx {
    fn encode(&self, v: &mut Vec<u8>) {
        self.encode_base(v);
        v.push_script_hash(&self.from);
        v.push_script_hash(&self.to);
        v.push_bytes(&self.script);
        v.push_asset(&self.amount);
        v.push_bytes(&self.memo);
    }
}

impl DecodeTx<TransferTx> for TransferTx {
    fn decode(cur: &mut Cursor<&[u8]>, tx: Tx) -> Option<TransferTx> {
        assert_eq!(tx.tx_type, TxType::TRANSFER);
        let from = cur.take_script_hash().ok()?;
        let to = cur.take_script_hash().ok()?;
        let script = Script::new(cur.take_bytes().ok()?);
        let amount = cur.take_asset().ok()?;
        let memo = cur.take_bytes().ok()?;
        Some(TransferTx {
            base: tx,
            from,
            to,
            script,
            amount,
            memo
        })
    }
}

tx_deref!(RewardTx);
tx_deref!(BondTx);
tx_deref!(TransferTx);

tx_sign!(BondTx);
tx_sign!(TransferTx);

#[cfg(test)]
mod tests {
    use ::std::str::FromStr;

    use crate::crypto;
    use super::*;

    macro_rules! cmp_base_tx {
        ($id:ident, $ty:expr, $ts:expr, $fee:expr) => {
            assert_eq!($id.tx_type, $ty);
            assert_eq!($id.timestamp, $ts);
            assert_eq!($id.fee.to_string(), $fee);
        }
    }

    #[test]
    fn test_encode_tx_with_sigs() {
        let to = crypto::KeyPair::gen_keypair();
        let reward_tx = TxVariant::RewardTx(RewardTx {
            base: Tx {
                tx_type: TxType::REWARD,
                timestamp: 123,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            to: to.0.into(),
            rewards: vec![get_asset("1.50 GOLD"), get_asset("1.0 SILVER")]
        });

        let mut v = vec![];
        reward_tx.encode_with_sigs(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        TxVariant::decode_with_sigs(&mut c).unwrap();
    }

    #[test]
    fn test_encode_reward() {
        let to = crypto::KeyPair::gen_keypair();
        let reward_tx = RewardTx {
            base: Tx {
                tx_type: TxType::REWARD,
                timestamp: 123,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            to: to.0.into(),
            rewards: vec![get_asset("1.50 GOLD"), get_asset("1.0 SILVER")]
        };

        let mut v = vec![];
        reward_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = RewardTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::REWARD, 123, "123 GOLD");
        assert_eq!(reward_tx.to, dec.to);
        assert_eq!(reward_tx.rewards.len(), dec.rewards.len());
        assert_eq!(reward_tx.rewards[0].to_string(), dec.rewards[0].to_string());
        assert_eq!(reward_tx.rewards[1].to_string(), dec.rewards[1].to_string());
    }

    #[test]
    fn test_encode_bond() {
        let minter = crypto::KeyPair::gen_keypair();
        let staker = crypto::KeyPair::gen_keypair();
        let bond_tx = BondTx {
            base: Tx {
                tx_type: TxType::BOND,
                timestamp: 1230,
                fee: get_asset("123 GOLD"),
                signature_pairs: vec![]
            },
            minter: minter.0,
            staker: staker.0.into(),
            stake_amt: get_asset("1.0456 GOLD"),
            bond_fee: get_asset("1.00000000 GOLD")
        };

        let mut v = vec![];
        bond_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = BondTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::BOND, 1230, "123 GOLD");
        assert_eq!(bond_tx.minter, dec.minter);
        assert_eq!(bond_tx.staker, dec.staker);
        assert_eq!(bond_tx.stake_amt.to_string(), dec.stake_amt.to_string());
        assert_eq!(bond_tx.bond_fee.to_string(), dec.bond_fee.to_string());
    }

    #[test]
    fn test_encode_transfer() {
        let from = crypto::KeyPair::gen_keypair();
        let to = crypto::KeyPair::gen_keypair();
        let transfer_tx = TransferTx {
            base: Tx {
                tx_type: TxType::TRANSFER,
                timestamp: 1234567890,
                fee: get_asset("1.23 GOLD"),
                signature_pairs: vec![]
            },
            from: from.0.into(),
            to: to.0.into(),
            script: Script::new(vec![1, 2, 3, 4]),
            amount: get_asset("1.0456 GOLD"),
            memo: Vec::from(String::from("Hello world!").as_bytes())
        };

        let mut v = vec![];
        transfer_tx.encode(&mut v);

        let mut c = Cursor::<&[u8]>::new(&v);
        let base = Tx::decode_base(&mut c).unwrap();
        let dec = TransferTx::decode(&mut c, base).unwrap();

        cmp_base_tx!(dec, TxType::TRANSFER, 1234567890, "1.23 GOLD");
        assert_eq!(transfer_tx.from, dec.from);
        assert_eq!(transfer_tx.to, dec.to);
        assert_eq!(transfer_tx.script, Script::new(vec![1, 2, 3, 4]));
        assert_eq!(transfer_tx.amount.to_string(), dec.amount.to_string());
        assert_eq!(transfer_tx.memo, dec.memo);
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}
<|endoftext|>


<commit_before>use join_to_string::join;

use libsyntax2::{
    File, TextUnit, TextRange,
    ast::{self, AstNode, AttrsOwner, TypeParamsOwner, NameOwner},
    SyntaxKind::{COMMA, WHITESPACE},
    SyntaxNodeRef,
    algo::{
        Direction, siblings,
        find_leaf_at_offset,
        find_covering_node,
        ancestors,
    },
};

use {EditBuilder, Edit, find_node_at_offset};

#[derive(Debug)]
pub struct LocalEdit {
    pub edit: Edit,
    pub cursor_position: Option<TextUnit>,
}

pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let syntax = file.syntax();

    let comma = find_leaf_at_offset(syntax, offset).find(|leaf| leaf.kind() == COMMA)?;
    let left = non_trivia_sibling(comma, Direction::Backward)?;
    let right = non_trivia_sibling(comma, Direction::Forward)?;
    Some(move || {
        let mut edit = EditBuilder::new();
        edit.replace(left.range(), right.text().to_string());
        edit.replace(right.range(), left.text().to_string());
        LocalEdit {
            edit: edit.finish(),
            cursor_position: None,
        }
    })
}

pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    Some(move || {
        let derive_attr = nominal
            .attrs()
            .filter_map(|x| x.as_call())
            .filter(|(name, _arg)| name == "derive")
            .map(|(_name, arg)| arg)
            .next();
        let mut edit = EditBuilder::new();
        let offset = match derive_attr {
            None => {
                let node_start = nominal.syntax().range().start();
                edit.insert(node_start, "#[derive()]\n".to_string());
                node_start + TextUnit::of_str("#[derive(")
            }
            Some(tt) => {
                tt.syntax().range().end() - TextUnit::of_char(')')
            }
        };
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn add_impl<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    let name = nominal.name()?;

    Some(move || {
        let type_params = nominal.type_param_list();
        let mut edit = EditBuilder::new();
        let start_offset = nominal.syntax().range().end();
        let mut buf = String::new();
        buf.push_str("\n\nimpl");
        if let Some(type_params) = type_params {
            type_params.syntax().text()
                .push_to(&mut buf);
        }
        buf.push_str(" ");
        buf.push_str(name.text().as_str());
        if let Some(type_params) = type_params {
            let lifetime_params = type_params.lifetime_params().filter_map(|it| it.lifetime()).map(|it| it.text());
            let type_params = type_params.type_params().filter_map(|it| it.name()).map(|it| it.text());
            join(lifetime_params.chain(type_params))
                .surround_with("<", ">")
                .to_buf(&mut buf);
        }
        buf.push_str(" {\n");
        let offset = start_offset + TextUnit::of_str(&buf);
        buf.push_str("\n}");
        edit.insert(start_offset, buf);
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn introduce_variable<'a>(file: &'a File, range: TextRange) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let node = find_covering_node(file.syntax(), range);
    let expr = ancestors(node).filter_map(ast::Expr::cast).next()?;
    let anchor_stmt = ancestors(expr.syntax()).filter_map(ast::Stmt::cast).next()?;
    let indent = anchor_stmt.syntax().prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    Some(move || {
        let mut buf = String::new();
        buf.push_str("let var_name = ");
        expr.syntax().text().push_to(&mut buf);
        buf.push_str(";");
        indent.text().push_to(&mut buf);

        let mut edit = EditBuilder::new();
        edit.replace(expr.syntax().range(), "var_name".to_string());
        edit.insert(anchor_stmt.syntax().range().start(), buf);
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(anchor_stmt.syntax().range().start() + TextUnit::of_str("let ")),
        }
    })
}

fn non_trivia_sibling(node: SyntaxNodeRef, direction: Direction) -> Option<SyntaxNodeRef> {
    siblings(node, direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, check_action_range};

    #[test]
    fn test_swap_comma() {
        check_action(
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
            |file, off| flip_comma(file, off).map(|f| f()),
        )
    }

    #[test]
    fn test_add_derive() {
        check_action(
            "struct Foo { a: i32, <|>}",
            "#[derive(<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo { <|> a: i32, }",
            "#[derive(<|>)]\nstruct Foo {  a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "#[derive(Clone)]\nstruct Foo { a: i32<|>, }",
            "#[derive(Clone<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_add_impl() {
        check_action(
            "struct Foo {<|>}\n",
            "struct Foo {}\n\nimpl Foo {\n<|>\n}\n",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<T: Clone> {<|>}",
            "struct Foo<T: Clone> {}\n\nimpl<T: Clone> Foo<T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<'a, T: Foo<'a>> {<|>}",
            "struct Foo<'a, T: Foo<'a>> {}\n\nimpl<'a, T: Foo<'a>> Foo<'a, T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var() {
        check_action_range(
            "
fn foo() {
    foo(<|>1 + 1<|>);
}", "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

}
<commit_msg>better introduce<commit_after>use join_to_string::join;

use libsyntax2::{
    File, TextUnit, TextRange,
    ast::{self, AstNode, AttrsOwner, TypeParamsOwner, NameOwner},
    SyntaxKind::{COMMA, WHITESPACE},
    SyntaxNodeRef,
    algo::{
        Direction, siblings,
        find_leaf_at_offset,
        find_covering_node,
        ancestors,
    },
};

use {EditBuilder, Edit, find_node_at_offset};

#[derive(Debug)]
pub struct LocalEdit {
    pub edit: Edit,
    pub cursor_position: Option<TextUnit>,
}

pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let syntax = file.syntax();

    let comma = find_leaf_at_offset(syntax, offset).find(|leaf| leaf.kind() == COMMA)?;
    let left = non_trivia_sibling(comma, Direction::Backward)?;
    let right = non_trivia_sibling(comma, Direction::Forward)?;
    Some(move || {
        let mut edit = EditBuilder::new();
        edit.replace(left.range(), right.text().to_string());
        edit.replace(right.range(), left.text().to_string());
        LocalEdit {
            edit: edit.finish(),
            cursor_position: None,
        }
    })
}

pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    Some(move || {
        let derive_attr = nominal
            .attrs()
            .filter_map(|x| x.as_call())
            .filter(|(name, _arg)| name == "derive")
            .map(|(_name, arg)| arg)
            .next();
        let mut edit = EditBuilder::new();
        let offset = match derive_attr {
            None => {
                let node_start = nominal.syntax().range().start();
                edit.insert(node_start, "#[derive()]\n".to_string());
                node_start + TextUnit::of_str("#[derive(")
            }
            Some(tt) => {
                tt.syntax().range().end() - TextUnit::of_char(')')
            }
        };
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn add_impl<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    let name = nominal.name()?;

    Some(move || {
        let type_params = nominal.type_param_list();
        let mut edit = EditBuilder::new();
        let start_offset = nominal.syntax().range().end();
        let mut buf = String::new();
        buf.push_str("\n\nimpl");
        if let Some(type_params) = type_params {
            type_params.syntax().text()
                .push_to(&mut buf);
        }
        buf.push_str(" ");
        buf.push_str(name.text().as_str());
        if let Some(type_params) = type_params {
            let lifetime_params = type_params.lifetime_params().filter_map(|it| it.lifetime()).map(|it| it.text());
            let type_params = type_params.type_params().filter_map(|it| it.name()).map(|it| it.text());
            join(lifetime_params.chain(type_params))
                .surround_with("<", ">")
                .to_buf(&mut buf);
        }
        buf.push_str(" {\n");
        let offset = start_offset + TextUnit::of_str(&buf);
        buf.push_str("\n}");
        edit.insert(start_offset, buf);
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn introduce_variable<'a>(file: &'a File, range: TextRange) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let node = find_covering_node(file.syntax(), range);
    let expr = ancestors(node).filter_map(ast::Expr::cast).next()?;
    let anchor_stmt = ancestors(expr.syntax()).filter_map(ast::Stmt::cast).next()?;
    let indent = anchor_stmt.syntax().prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    Some(move || {
        let mut buf = String::new();
        let mut edit = EditBuilder::new();

        buf.push_str("let var_name = ");
        expr.syntax().text().push_to(&mut buf);
        if expr.syntax().range().start() == anchor_stmt.syntax().range().start() {
            edit.replace(expr.syntax().range(), buf);
        } else {
            buf.push_str(";");
            indent.text().push_to(&mut buf);
            edit.replace(expr.syntax().range(), "var_name".to_string());
            edit.insert(anchor_stmt.syntax().range().start(), buf);
        }
        let cursor_position = anchor_stmt.syntax().range().start() + TextUnit::of_str("let ");
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(cursor_position),
        }
    })
}

fn non_trivia_sibling(node: SyntaxNodeRef, direction: Direction) -> Option<SyntaxNodeRef> {
    siblings(node, direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, check_action_range};

    #[test]
    fn test_swap_comma() {
        check_action(
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
            |file, off| flip_comma(file, off).map(|f| f()),
        )
    }

    #[test]
    fn test_add_derive() {
        check_action(
            "struct Foo { a: i32, <|>}",
            "#[derive(<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo { <|> a: i32, }",
            "#[derive(<|>)]\nstruct Foo {  a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "#[derive(Clone)]\nstruct Foo { a: i32<|>, }",
            "#[derive(Clone<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_add_impl() {
        check_action(
            "struct Foo {<|>}\n",
            "struct Foo {}\n\nimpl Foo {\n<|>\n}\n",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<T: Clone> {<|>}",
            "struct Foo<T: Clone> {}\n\nimpl<T: Clone> Foo<T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<'a, T: Foo<'a>> {<|>}",
            "struct Foo<'a, T: Foo<'a>> {}\n\nimpl<'a, T: Foo<'a>> Foo<'a, T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var_simple() {
        check_action_range(
            "
fn foo() {
    foo(<|>1 + 1<|>);
}", "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }
    #[test]
    fn test_intrdoduce_var_expr_stmt() {
check_action_range(
            "
fn foo() {
    <|>1 + 1<|>;
}", "
fn foo() {
    let <|>var_name = 1 + 1;
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

}
<|endoftext|>


<commit_before>use ra_proc_macro::msg::{self, Message};
use ra_proc_macro_srv::{expand_task, list_macros};

use std::io;

fn read_request() -> Result<Option<msg::Request>, io::Error> {
    let stdin = io::stdin();
    let mut stdin = stdin.lock();
    msg::Request::read(&mut stdin)
}

fn write_response(res: Result<msg::Response, String>) -> Result<(), io::Error> {
    let msg: msg::Response = match res {
        Ok(res) => res,
        Err(err) => msg::Response::Error(msg::ResponseError {
            code: msg::ErrorCode::ExpansionError,
            message: err,
        }),
    };

    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    msg.write(&mut stdout)
}
fn main() {
    loop {
        let req = match read_request() {
            Err(err) => {
                eprintln!("Read message error on ra_proc_macro_srv: {}", err.to_string());
                continue;
            }
            Ok(None) => continue,
            Ok(Some(req)) => req,
        };

        match req {
            msg::Request::ListMacro(task) => {
                if let Err(err) =
                    write_response(list_macros(&task).map(|it| msg::Response::ListMacro(it)))
                {
                    eprintln!("Write message error on list macro: {}", err);
                }
            }
            msg::Request::ExpansionMacro(task) => {
                if let Err(err) =
                    write_response(expand_task(&task).map(|it| msg::Response::ExpansionMacro(it)))
                {
                    eprintln!("Write message error on expansion macro: {}", err);
                }
            }
        }
    }
}
<commit_msg>Add doc comment on main.rs<commit_after>//! Driver for proc macro server

use ra_proc_macro::msg::{self, Message};
use ra_proc_macro_srv::{expand_task, list_macros};

use std::io;

fn read_request() -> Result<Option<msg::Request>, io::Error> {
    let stdin = io::stdin();
    let mut stdin = stdin.lock();
    msg::Request::read(&mut stdin)
}

fn write_response(res: Result<msg::Response, String>) -> Result<(), io::Error> {
    let msg: msg::Response = match res {
        Ok(res) => res,
        Err(err) => msg::Response::Error(msg::ResponseError {
            code: msg::ErrorCode::ExpansionError,
            message: err,
        }),
    };

    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    msg.write(&mut stdout)
}
fn main() {
    loop {
        let req = match read_request() {
            Err(err) => {
                eprintln!("Read message error on ra_proc_macro_srv: {}", err.to_string());
                continue;
            }
            Ok(None) => continue,
            Ok(Some(req)) => req,
        };

        match req {
            msg::Request::ListMacro(task) => {
                if let Err(err) =
                    write_response(list_macros(&task).map(|it| msg::Response::ListMacro(it)))
                {
                    eprintln!("Write message error on list macro: {}", err);
                }
            }
            msg::Request::ExpansionMacro(task) => {
                if let Err(err) =
                    write_response(expand_task(&task).map(|it| msg::Response::ExpansionMacro(it)))
                {
                    eprintln!("Write message error on expansion macro: {}", err);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Configuration for setting up database pools
//!
//! - `DATABASE_URL`: The URL of the postgres database to use.
//! - `READ_ONLY_REPLICA_URL`: The URL of an optional postgres read-only replica database.
//! - `DB_PRIMARY_POOL_SIZE`: The number of connections of the primary database.
//! - `DB_REPLICA_POOL_SIZE`: The number of connections of the read-only / replica database.
//! - `DB_OFFLINE`: If set to `leader` then use the read-only follower as if it was the leader.
//!   If set to `follower` then act as if `READ_ONLY_REPLICA_URL` was unset.
//! - `READ_ONLY_MODE`: If defined (even as empty) then force all connections to be read-only.

use crate::env;

pub struct DatabasePools {
    /// Settings for the primary database. This is usually writeable, but will be read-only in
    /// some configurations.
    pub primary: DbPoolConfig,
    /// An optional follower database. Always read-only.
    pub replica: Option<DbPoolConfig>,
}

#[derive(Debug)]
pub struct DbPoolConfig {
    pub url: String,
    pub read_only_mode: bool,
    pub pool_size: u32,
    pub min_idle: Option<u32>,
}

impl DatabasePools {
    pub fn are_all_read_only(&self) -> bool {
        self.primary.read_only_mode
    }

    pub fn primary_max_pool_size(&self) -> u32 {
        self.primary.pool_size
    }

    pub fn replica_max_pool_size(&self) -> u32 {
        self.replica
            .as_ref()
            .map(|config| config.pool_size)
            .unwrap_or(Self::DEFAULT_POOL_SIZE)
    }
}

impl DatabasePools {
    const DEFAULT_POOL_SIZE: u32 = 3;

    /// Load settings for one or more database pools from the environment
    ///
    /// # Panics
    ///
    /// This function panics if `DB_OFFLINE=leader` but `READ_ONLY_REPLICA_URL` is unset.
    pub fn full_from_environment() -> Self {
        let leader_url = env("DATABASE_URL");
        let follower_url = dotenv::var("READ_ONLY_REPLICA_URL").ok();
        let read_only_mode = dotenv::var("READ_ONLY_MODE").is_ok();

        let primary_pool_size = match dotenv::var("DB_PRIMARY_POOL_SIZE") {
            Ok(num) => num.parse().expect("couldn't parse DB_PRIMARY_POOL_SIZE"),
            _ => Self::DEFAULT_POOL_SIZE,
        };

        let replica_pool_size = match dotenv::var("DB_REPLICA_POOL_SIZE") {
            Ok(num) => num.parse().expect("couldn't parse DB_REPLICA_POOL_SIZE"),
            _ => Self::DEFAULT_POOL_SIZE,
        };

        let primary_min_idle = match dotenv::var("DB_PRIMARY_MIN_IDLE") {
            Ok(num) => Some(num.parse().expect("couldn't parse DB_PRIMARY_MIN_IDLE")),
            _ => None,
        };

        let replica_min_idle = match dotenv::var("DB_REPLICA_MIN_IDLE") {
            Ok(num) => Some(num.parse().expect("couldn't parse DB_REPLICA_MIN_IDLE")),
            _ => None,
        };

        match dotenv::var("DB_OFFLINE").as_deref() {
            // The actual leader is down, use the follower in read-only mode as the primary and
            // don't configure a replica.
            Ok("leader") => Self {
                primary: DbPoolConfig {
                    url: follower_url
                        .expect("Must set `READ_ONLY_REPLICA_URL` when using `DB_OFFLINE=leader`."),
                    read_only_mode: true,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: None,
            },
            // The follower is down, don't configure the replica.
            Ok("follower") => Self {
                primary: DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: None,
            },
            _ => Self {
                primary: DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: follower_url.map(|url| DbPoolConfig {
                    url,
                    // Always enable read-only mode for the follower. In staging, we attach the
                    // same leader database to both environment variables and this ensures the
                    // connection is opened read-only even when attached to a writeable database.
                    read_only_mode: true,
                    pool_size: replica_pool_size,
                    min_idle: replica_min_idle,
                }),
            },
        }
    }

    pub fn test_from_environment() -> Self {
        DatabasePools {
            primary: DbPoolConfig {
                url: env("TEST_DATABASE_URL"),
                read_only_mode: false,
                pool_size: 1,
                min_idle: None,
            },
            replica: None,
        }
    }
}
<commit_msg>Add description of env vars<commit_after>//! Configuration for setting up database pools
//!
//! - `DATABASE_URL`: The URL of the postgres database to use.
//! - `READ_ONLY_REPLICA_URL`: The URL of an optional postgres read-only replica database.
//! - `DB_PRIMARY_POOL_SIZE`: The number of connections of the primary database.
//! - `DB_REPLICA_POOL_SIZE`: The number of connections of the read-only / replica database.
//! - `DB_PRIMARY_MIN_IDLE`: The primary pool will maintain at least this number of connections.
//! - `DB_REPLICA_MIN_IDLE`: The replica pool will maintain at least this number of connections.
//! - `DB_OFFLINE`: If set to `leader` then use the read-only follower as if it was the leader.
//!   If set to `follower` then act as if `READ_ONLY_REPLICA_URL` was unset.
//! - `READ_ONLY_MODE`: If defined (even as empty) then force all connections to be read-only.

use crate::env;

pub struct DatabasePools {
    /// Settings for the primary database. This is usually writeable, but will be read-only in
    /// some configurations.
    pub primary: DbPoolConfig,
    /// An optional follower database. Always read-only.
    pub replica: Option<DbPoolConfig>,
}

#[derive(Debug)]
pub struct DbPoolConfig {
    pub url: String,
    pub read_only_mode: bool,
    pub pool_size: u32,
    pub min_idle: Option<u32>,
}

impl DatabasePools {
    pub fn are_all_read_only(&self) -> bool {
        self.primary.read_only_mode
    }

    pub fn primary_max_pool_size(&self) -> u32 {
        self.primary.pool_size
    }

    pub fn replica_max_pool_size(&self) -> u32 {
        self.replica
            .as_ref()
            .map(|config| config.pool_size)
            .unwrap_or(Self::DEFAULT_POOL_SIZE)
    }
}

impl DatabasePools {
    const DEFAULT_POOL_SIZE: u32 = 3;

    /// Load settings for one or more database pools from the environment
    ///
    /// # Panics
    ///
    /// This function panics if `DB_OFFLINE=leader` but `READ_ONLY_REPLICA_URL` is unset.
    pub fn full_from_environment() -> Self {
        let leader_url = env("DATABASE_URL");
        let follower_url = dotenv::var("READ_ONLY_REPLICA_URL").ok();
        let read_only_mode = dotenv::var("READ_ONLY_MODE").is_ok();

        let primary_pool_size = match dotenv::var("DB_PRIMARY_POOL_SIZE") {
            Ok(num) => num.parse().expect("couldn't parse DB_PRIMARY_POOL_SIZE"),
            _ => Self::DEFAULT_POOL_SIZE,
        };

        let replica_pool_size = match dotenv::var("DB_REPLICA_POOL_SIZE") {
            Ok(num) => num.parse().expect("couldn't parse DB_REPLICA_POOL_SIZE"),
            _ => Self::DEFAULT_POOL_SIZE,
        };

        let primary_min_idle = match dotenv::var("DB_PRIMARY_MIN_IDLE") {
            Ok(num) => Some(num.parse().expect("couldn't parse DB_PRIMARY_MIN_IDLE")),
            _ => None,
        };

        let replica_min_idle = match dotenv::var("DB_REPLICA_MIN_IDLE") {
            Ok(num) => Some(num.parse().expect("couldn't parse DB_REPLICA_MIN_IDLE")),
            _ => None,
        };

        match dotenv::var("DB_OFFLINE").as_deref() {
            // The actual leader is down, use the follower in read-only mode as the primary and
            // don't configure a replica.
            Ok("leader") => Self {
                primary: DbPoolConfig {
                    url: follower_url
                        .expect("Must set `READ_ONLY_REPLICA_URL` when using `DB_OFFLINE=leader`."),
                    read_only_mode: true,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: None,
            },
            // The follower is down, don't configure the replica.
            Ok("follower") => Self {
                primary: DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: None,
            },
            _ => Self {
                primary: DbPoolConfig {
                    url: leader_url,
                    read_only_mode,
                    pool_size: primary_pool_size,
                    min_idle: primary_min_idle,
                },
                replica: follower_url.map(|url| DbPoolConfig {
                    url,
                    // Always enable read-only mode for the follower. In staging, we attach the
                    // same leader database to both environment variables and this ensures the
                    // connection is opened read-only even when attached to a writeable database.
                    read_only_mode: true,
                    pool_size: replica_pool_size,
                    min_idle: replica_min_idle,
                }),
            },
        }
    }

    pub fn test_from_environment() -> Self {
        DatabasePools {
            primary: DbPoolConfig {
                url: env("TEST_DATABASE_URL"),
                read_only_mode: false,
                pool_size: 1,
                min_idle: None,
            },
            replica: None,
        }
    }
}
<|endoftext|>


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

//! The implementations of the `Uniform` distribution for integer types.

use core::mem;

use {Rng};
use distributions::{Distribution, Uniform};

impl Distribution<isize> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> isize {
        if mem::size_of::<isize>() == 4 {
            rng.gen::<i32>() as isize
        } else {
            rng.gen::<i64>() as isize
        }
    }
}

impl Distribution<i8> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i8 {
        rng.next_u32() as i8
    }
}

impl Distribution<i16> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i16 {
        rng.next_u32() as i16
    }
}

impl Distribution<i32> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i32 {
        rng.next_u32() as i32
    }
}

impl Distribution<i64> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i64 {
        rng.next_u64() as i64
    }
}

#[cfg(feature = "i128_support")]
impl Distribution<i128> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i128 {
        rng.gen::<u128>() as i128
    }
}

impl Distribution<usize> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        if mem::size_of::<usize>() == 4 {
            rng.gen::<u32>() as usize
        } else {
            rng.gen::<u64>() as usize
        }
    }
}

impl Distribution<u8> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
        rng.next_u32() as u8
    }
}

impl Distribution<u16> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u16 {
        rng.next_u32() as u16
    }
}

impl Distribution<u32> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u32 {
        rng.next_u32()
    }
}

impl Distribution<u64> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        rng.next_u64()
    }
}

#[cfg(feature = "i128_support")]
impl Distribution<u128> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u128 {
        ((rng.next_u64() as u128) << 64) | (rng.next_u64() as u128)
    }
}

#[cfg(test)]
mod tests {
    use Rng;
    use distributions::{Uniform};

    #[test]
    fn test_integers() {
        let mut rng = ::test::rng(806);

        rng.sample::<isize, _>(Uniform);
        rng.sample::<i8, _>(Uniform);
        rng.sample::<i16, _>(Uniform);
        rng.sample::<i32, _>(Uniform);
        rng.sample::<i64, _>(Uniform);
        #[cfg(feature = "i128_support")]
        rng.sample::<i128, _>(Uniform);

        rng.sample::<usize, _>(Uniform);
        rng.sample::<u8, _>(Uniform);
        rng.sample::<u16, _>(Uniform);
        rng.sample::<u32, _>(Uniform);
        rng.sample::<u64, _>(Uniform);
        #[cfg(feature = "i128_support")]
        rng.sample::<u128, _>(Uniform);
    }
}
<commit_msg>Don't assume usize is 4 or 8 bytes<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The implementations of the `Uniform` distribution for integer types.

use {Rng};
use distributions::{Distribution, Uniform};

impl Distribution<isize> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> isize {
        rng.gen::<usize>() as isize
    }
}

impl Distribution<i8> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i8 {
        rng.next_u32() as i8
    }
}

impl Distribution<i16> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i16 {
        rng.next_u32() as i16
    }
}

impl Distribution<i32> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i32 {
        rng.next_u32() as i32
    }
}

impl Distribution<i64> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i64 {
        rng.next_u64() as i64
    }
}

#[cfg(feature = "i128_support")]
impl Distribution<i128> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i128 {
        rng.gen::<u128>() as i128
    }
}

impl Distribution<usize> for Uniform {
    #[inline]
    #[cfg(any(target_pointer_width = "32", target_pointer_width = "16"))]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        rng.next_u32() as usize
    }

    #[inline]
    #[cfg(target_pointer_width = "64")]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
        rng.next_u64() as usize
    }
}

impl Distribution<u8> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u8 {
        rng.next_u32() as u8
    }
}

impl Distribution<u16> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u16 {
        rng.next_u32() as u16
    }
}

impl Distribution<u32> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u32 {
        rng.next_u32()
    }
}

impl Distribution<u64> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        rng.next_u64()
    }
}

#[cfg(feature = "i128_support")]
impl Distribution<u128> for Uniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u128 {
        ((rng.next_u64() as u128) << 64) | (rng.next_u64() as u128)
    }
}

#[cfg(test)]
mod tests {
    use Rng;
    use distributions::{Uniform};

    #[test]
    fn test_integers() {
        let mut rng = ::test::rng(806);

        rng.sample::<isize, _>(Uniform);
        rng.sample::<i8, _>(Uniform);
        rng.sample::<i16, _>(Uniform);
        rng.sample::<i32, _>(Uniform);
        rng.sample::<i64, _>(Uniform);
        #[cfg(feature = "i128_support")]
        rng.sample::<i128, _>(Uniform);

        rng.sample::<usize, _>(Uniform);
        rng.sample::<u8, _>(Uniform);
        rng.sample::<u16, _>(Uniform);
        rng.sample::<u32, _>(Uniform);
        rng.sample::<u64, _>(Uniform);
        #[cfg(feature = "i128_support")]
        rng.sample::<u128, _>(Uniform);
    }
}
<|endoftext|>


<commit_before>//! Pointers and their events

use types::input_device::InputDevice;

use wlroots_sys::{wlr_button_state, wlr_event_pointer_axis, wlr_event_pointer_button,
                  wlr_event_pointer_motion, wlr_event_pointer_motion_absolute};

// TODO FIXME Document this shit man

#[derive(Debug)]
pub struct AxisEvent {
    event: *mut wlr_event_pointer_axis,
    device: InputDevice
}

#[derive(Debug)]
pub struct ButtonEvent {
    event: *mut wlr_event_pointer_button,
    device: InputDevice
}

#[derive(Debug)]
pub struct MotionEvent {
    event: *mut wlr_event_pointer_motion,
    device: InputDevice
}

#[derive(Debug)]
pub struct AbsoluteMotionEvent {
    event: *mut wlr_event_pointer_motion_absolute,
    device: InputDevice
}

impl ButtonEvent {
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_button) -> Self {
        ButtonEvent { device: InputDevice::from_ptr((*event).device),
                      event }
    }

    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    pub fn state(&self) -> wlr_button_state {
        unsafe { (*self.event).state }
    }

    pub fn button(&self) -> u32 {
        unsafe { (*self.event).button }
    }
}

impl AxisEvent {
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_axis) -> Self {
        AxisEvent { device: InputDevice::from_ptr((*event).device),
                    event }
    }

    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    pub fn delta(&self) -> f64 {
        unsafe { (*self.event).delta }
    }
}

impl MotionEvent {
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_motion) -> Self {
        MotionEvent { device: InputDevice::from_ptr((*event).device),
                      event }
    }

    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    pub fn delta(&self) -> (f64, f64) {
        unsafe { ((*self.event).delta_x, (*self.event).delta_y) }
    }
}

impl AbsoluteMotionEvent {
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_motion_absolute) -> Self {
        AbsoluteMotionEvent { device: InputDevice::from_ptr((*event).device),
                              event }
    }

    pub fn device(&self) -> &InputDevice {
        &self.device
    }
}
<commit_msg>Document the pointer events<commit_after>//! Pointers and their events

use types::input_device::InputDevice;

use wlroots_sys::{wlr_button_state, wlr_event_pointer_axis, wlr_event_pointer_button,
                  wlr_event_pointer_motion, wlr_event_pointer_motion_absolute};

#[derive(Debug)]
/// Event that triggers when the pointer device scrolls (e.g using a wheel
// or in the case of a touchpad when you use two fingers to scroll).
pub struct AxisEvent {
    event: *mut wlr_event_pointer_axis,
    device: InputDevice
}

#[derive(Debug)]
/// Event that triggers when a button is pressed (e.g left click, right click,
/// a gaming mouse button, etc.).
pub struct ButtonEvent {
    event: *mut wlr_event_pointer_button,
    device: InputDevice
}

#[derive(Debug)]
/// Event that triggers when the pointer moves.
pub struct MotionEvent {
    event: *mut wlr_event_pointer_motion,
    device: InputDevice
}

#[derive(Debug)]
/// Event that triggers when data from a device that supports absolute motion
/// sends data to the compositor.
///
/// For more information on absolute motion, [see this link](https://wayland.freedesktop.org/libinput/doc/latest/absolute_axes.html).
pub struct AbsoluteMotionEvent {
    event: *mut wlr_event_pointer_motion_absolute,
    device: InputDevice
}

impl ButtonEvent {
    /// Constructs a `ButtonEvent` from the raw event pointer.
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_button) -> Self {
        ButtonEvent { device: InputDevice::from_ptr((*event).device),
                      event }
    }

    /// Get the device this event refers to.
    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    /// Get the state of the button (e.g pressed or released).
    pub fn state(&self) -> wlr_button_state {
        unsafe { (*self.event).state }
    }

    /// Get the value of the button pressed. This will generally be an atomically
    /// increasing value, with e.g left click being 1 and right click being 2...
    ///
    /// We make no guarantees that 1 always maps to left click, as this is device
    /// driver specific.
    pub fn button(&self) -> u32 {
        unsafe { (*self.event).button }
    }
}

impl AxisEvent {
    /// Constructs a `AxisEvent` from a raw event pointer.
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_axis) -> Self {
        AxisEvent { device: InputDevice::from_ptr((*event).device),
                    event }
    }

    /// Get the device this event refers to.
    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    /// Get the change from the last axis value.
    ///
    /// Useful to determine e.g how much to scroll.
    pub fn delta(&self) -> f64 {
        unsafe { (*self.event).delta }
    }
}

impl MotionEvent {
    /// Constructs a `MotionEvent` from a raw event pointer.
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_motion) -> Self {
        MotionEvent { device: InputDevice::from_ptr((*event).device),
                      event }
    }

    /// Get the device this event refers to.
    pub fn device(&self) -> &InputDevice {
        &self.device
    }

    /// Get the change from the last positional value.
    ///
    /// Returned in (x, y) form.
    ///
    /// Note you should not cast this to a type with less precision,
    /// otherwise you'll lose important motion data which can cause bugs
    /// (e.g see [this fun wlc bug](https://github.com/Cloudef/wlc/issues/181)).
    pub fn delta(&self) -> (f64, f64) {
        unsafe { ((*self.event).delta_x, (*self.event).delta_y) }
    }
}

impl AbsoluteMotionEvent {
    /// Construct an `AbsoluteMotionEvent` from a raw event pointer.
    pub unsafe fn from_ptr(event: *mut wlr_event_pointer_motion_absolute) -> Self {
        AbsoluteMotionEvent { device: InputDevice::from_ptr((*event).device),
                              event }
    }

    /// Get the device this event refers to.
    pub fn device(&self) -> &InputDevice {
        &self.device
    }
}
<|endoftext|>


<commit_before>
macro_rules! yaml_tuple2 {
    ($a:ident, $v:ident, $c:ident) => {{
            if let Some(vec) = $v.as_vec() {
                for ys in vec {
                    if let Some(tup) = ys.as_vec() {
                        debug_assert_eq!(2, tup.len());
                        $a = $a.$c(yaml_str!(tup[0]), yaml_str!(tup[1]));
                    } else {
                        panic!("Failed to convert YAML value to vec");
                    }
                }
            } else {
                panic!("Failed to convert YAML value to vec");
            }
            $a
        }
    };
}

macro_rules! yaml_tuple3 {
    ($a:ident, $v:ident, $c:ident) => {{
            if let Some(vec) = $v.as_vec() {
                for ys in vec {
                    if let Some(tup) = ys.as_vec() {
                        debug_assert_eq!(3, tup.len());
                        $a = $a.$c(yaml_str!(tup[0]), yaml_opt_str!(tup[1]), yaml_str!(tup[2]));
                    } else {
                        panic!("Failed to convert YAML value to vec");
                    }
                }
            } else {
                panic!("Failed to convert YAML value to vec");
            }
            $a
        }
    };
}

macro_rules! yaml_vec_or_str {
    ($v:ident, $a:ident, $c:ident) => {{
            let maybe_vec = $v.as_vec();
            if let Some(vec) = maybe_vec {
                for ys in vec {
                    if let Some(s) = ys.as_str() {
                        $a = $a.$c(s);
                    } else {
                        panic!("Failed to convert YAML value {:?} to a string", ys);
                    }
                }
            } else {
                if let Some(s) = $v.as_str() {
                    $a = $a.$c(s);
                } else {
                    panic!("Failed to convert YAML value {:?} to either a vec or string", $v);
                }
            }
            $a
        }
    };
}

macro_rules! yaml_opt_str {
    ($v:expr) => {{
        if $v.is_null() {
            Some($v.as_str().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)))
        } else {
            None
        }
    }};
}

macro_rules! yaml_str {
    ($v:expr) => {{
        $v.as_str().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v))
    }};
}

macro_rules! yaml_to_str {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c(yaml_str!($v))
    }};
}

macro_rules! yaml_to_bool {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_bool().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)))
    }};
}

macro_rules! yaml_to_u64 {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_i64().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)) as u64)
    }};
}

macro_rules! yaml_to_usize {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_i64().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)) as usize)
    }};
}
<commit_msg>fix unused_macros warnings for non-yaml build<commit_after>#[cfg(feature = "yaml")]
macro_rules! yaml_tuple2 {
    ($a:ident, $v:ident, $c:ident) => {{
            if let Some(vec) = $v.as_vec() {
                for ys in vec {
                    if let Some(tup) = ys.as_vec() {
                        debug_assert_eq!(2, tup.len());
                        $a = $a.$c(yaml_str!(tup[0]), yaml_str!(tup[1]));
                    } else {
                        panic!("Failed to convert YAML value to vec");
                    }
                }
            } else {
                panic!("Failed to convert YAML value to vec");
            }
            $a
        }
    };
}

#[cfg(feature = "yaml")]
macro_rules! yaml_tuple3 {
    ($a:ident, $v:ident, $c:ident) => {{
            if let Some(vec) = $v.as_vec() {
                for ys in vec {
                    if let Some(tup) = ys.as_vec() {
                        debug_assert_eq!(3, tup.len());
                        $a = $a.$c(yaml_str!(tup[0]), yaml_opt_str!(tup[1]), yaml_str!(tup[2]));
                    } else {
                        panic!("Failed to convert YAML value to vec");
                    }
                }
            } else {
                panic!("Failed to convert YAML value to vec");
            }
            $a
        }
    };
}

#[cfg(feature = "yaml")]
macro_rules! yaml_vec_or_str {
    ($v:ident, $a:ident, $c:ident) => {{
            let maybe_vec = $v.as_vec();
            if let Some(vec) = maybe_vec {
                for ys in vec {
                    if let Some(s) = ys.as_str() {
                        $a = $a.$c(s);
                    } else {
                        panic!("Failed to convert YAML value {:?} to a string", ys);
                    }
                }
            } else {
                if let Some(s) = $v.as_str() {
                    $a = $a.$c(s);
                } else {
                    panic!("Failed to convert YAML value {:?} to either a vec or string", $v);
                }
            }
            $a
        }
    };
}

#[cfg(feature = "yaml")]
macro_rules! yaml_opt_str {
    ($v:expr) => {{
        if $v.is_null() {
            Some($v.as_str().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)))
        } else {
            None
        }
    }};
}

#[cfg(feature = "yaml")]
macro_rules! yaml_str {
    ($v:expr) => {{
        $v.as_str().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v))
    }};
}

#[cfg(feature = "yaml")]
macro_rules! yaml_to_str {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c(yaml_str!($v))
    }};
}

#[cfg(feature = "yaml")]
macro_rules! yaml_to_bool {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_bool().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)))
    }};
}

#[cfg(feature = "yaml")]
macro_rules! yaml_to_u64 {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_i64().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)) as u64)
    }};
}

#[cfg(feature = "yaml")]
macro_rules! yaml_to_usize {
    ($a:ident, $v:ident, $c:ident) => {{
        $a.$c($v.as_i64().unwrap_or_else(|| panic!("failed to convert YAML {:?} value to a string", $v)) as usize)
    }};
}
<|endoftext|>