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

use {Intrinsic, i, f, v};
use rustc::middle::ty;

macro_rules! p {
    ($name: expr, ($($inputs: tt),*) -> $output: tt) => {
        plain!(concat!("llvm.x86.", $name), ($($inputs),*) -> $output)
    }
}

pub fn find<'tcx>(_tcx: &ty::ctxt<'tcx>, name: &str) -> Option<Intrinsic> {
    if name.starts_with("mm_") {
        Some(match &name["mm_".len()..] {
            "sqrt_ps" => plain!("llvm.sqrt.v4f32", (f32x4) -> f32x4),
            "sqrt_pd" => plain!("llvm.sqrt.v2f64", (f64x2) -> f64x2),

            "movemask_ps" => p!("sse.movmsk.ps", (f32x4) -> i32),
            "max_ps" => p!("sse.max.ps", (f32x4, f32x4) -> f32x4),
            "min_ps" => p!("sse.min.ps", (f32x4, f32x4) -> f32x4),
            "rsqrt_ps" => p!("sse.rsqrt.ps", (f32x4) -> f32x4),
            "rcp_ps" => p!("sse.rcp.ps", (f32x4) -> f32x4),

            "adds_epi16" => p!("sse2.padds.w", (i16x8, i16x8) -> i16x8),
            "adds_epi8" => p!("sse2.padds.b", (i8x16, i8x16) -> i8x16),
            "adds_epu16" => p!("sse2.paddus.w", (i16x8, i16x8) -> i16x8),
            "adds_epu8" => p!("sse2.paddus.b", (i8x16, i8x16) -> i8x16),
            "avg_epu16" => p!("sse2.pavg.w", (i16x8, i16x8) -> i16x8),
            "avg_epu8" => p!("sse2.pavg.b", (i8x16, i8x16) -> i8x16),
            "madd_epi16" => p!("sse2.pmadd.wd", (i16x8, i16x8) -> i32x4),
            "max_epi16" => p!("sse2.pmaxs.w", (i16x8, i16x8) -> i16x8),
            "max_epu8" => p!("sse2.pmaxu.b", (i8x16, i8x16) -> i8x16),
            "max_pd" => p!("sse2.max.pd", (f64x2, f64x2) -> f64x2),
            "min_epi16" => p!("sse2.pmins.w", (i16x8, i16x8) -> i16x8),
            "min_epu8" => p!("sse2.pminu.b", (i8x16, i8x16) -> i8x16),
            "min_pd" => p!("sse2.min.pd", (f64x2, f64x2) -> f64x2),
            "movemask_pd" => p!("sse2.movmsk.pd", (f64x2) -> i32),
            "movemask_epi8" => p!("sse2.pmovmskb.128", (i8x16) -> i32),
            "mul_epu32" => p!("sse2.pmulu.dq", (i32x4, i32x4) -> i64x2),
            "mulhi_epi16" => p!("sse2.pmulh.w", (i8x16, i8x16) -> i8x16),
            "mulhi_epu16" => p!("sse2.pmulhu.w", (i8x16, i8x16) -> i8x16),
            "packs_epi16" => p!("sse2.packsswb.128", (i16x8, i16x8) -> i8x16),
            "packs_epi32" => p!("sse2.packssdw.128", (i32x4, i32x4) -> i16x8),
            "packus_epi16" => p!("sse2.packuswb.128", (i16x8, i16x8) -> i8x16),
            "sad_epu8" => p!("sse2.psad.bw", (i8x16, i8x16) -> i64x2),
            "subs_epi16" => p!("sse2.psubs.w", (i16x8, i16x8) -> i16x8),
            "subs_epi8" => p!("sse2.psubs.b", (i8x16, i8x16) -> i8x16),
            "subs_epu16" => p!("sse2.psubus.w", (i16x8, i16x8) -> i16x8),
            "subs_epu8" => p!("sse2.psubus.b", (i8x16, i8x16) -> i8x16),

            "addsub_pd" => p!("sse3.addsub.pd", (f64x2, f64x2) -> f64x2),
            "addsub_ps" => p!("sse3.addsub.ps", (f32x4, f32x4) -> f32x4),
            "hadd_pd" => p!("sse3.hadd.pd", (f64x2, f64x2) -> f64x2),
            "hadd_ps" => p!("sse3.hadd.ps", (f32x4, f32x4) -> f32x4),
            "hsub_pd" => p!("sse3.hsub.pd", (f64x2, f64x2) -> f64x2),
            "hsub_ps" => p!("sse3.hsub.ps", (f32x4, f32x4) -> f32x4),

            "abs_epi16" => p!("ssse3.pabs.w.128", (i16x8) -> i16x8),
            "abs_epi32" => p!("ssse3.pabs.d.128", (i32x4) -> i32x4),
            "abs_epi8" => p!("ssse3.pabs.b.128", (i8x16) -> i8x16),
            "hadd_epi16" => p!("ssse3.phadd.w.128", (i16x8, i16x8) -> i16x8),
            "hadd_epi32" => p!("ssse3.phadd.d.128", (i32x4, i32x4) -> i32x4),
            "hadds_epi16" => p!("ssse3.phadd.sw.128", (i16x8, i16x8) -> i16x8),
            "hsub_epi16" => p!("ssse3.phsub.w.128", (i16x8, i16x8) -> i16x8),
            "hsub_epi32" => p!("ssse3.phsub.d.128", (i32x4, i32x4) -> i32x4),
            "hsubs_epi16" => p!("ssse3.phsub.sw.128", (i16x8, i16x8) -> i16x8),
            "maddubs_epi16" => p!("ssse3.pmadd.ub.sw.128", (i8x16, i8x16) -> i16x8),
            "mulhrs_epi16" => p!("ssse3.pmul.hr.sw.128", (i16x8, i16x8) -> i16x8),
            "shuffle_epi8" => p!("ssse3.pshuf.b.128", (i8x16, i8x16) -> i8x16),
            "sign_epi16" => p!("ssse3.psign.w.128", (i16x8, i16x8) -> i16x8),
            "sign_epi32" => p!("ssse3.psign.d.128", (i32x4, i32x4) -> i32x4),
            "sign_epi8" => p!("ssse3.psign.b.128", (i8x16, i8x16) -> i8x16),

            "max_epi32" => p!("sse41.pmaxsd", (i32x4, i32x4) -> i32x4),
            "max_epi8" => p!("sse41.pmaxsb", (i8x16, i8x16) -> i8x16),
            "max_epu16" => p!("sse41.pmaxuw", (i16x8, i16x8) -> i16x8),
            "max_epu32" => p!("sse41.pmaxud", (i32x4, i32x4) -> i32x4),
            "min_epi32" => p!("sse41.pminsd", (i32x4, i32x4) -> i32x4),
            "min_epi8" => p!("sse41.pminsb", (i8x16, i8x16) -> i8x16),
            "min_epu16" => p!("sse41.pminuw", (i16x8, i16x8) -> i16x8),
            "min_epu32" => p!("sse41.pminud", (i32x4, i32x4) -> i32x4),
            "minpos_epu16" => p!("sse41.phminposuw", (i16x8) -> i16x8),
            "mul_epi32" => p!("sse41.muldq", (i32x4, i32x4) -> i64x2),
            "packus_epi32" => p!("sse41.packusdw", (i32x4, i32x4) -> i16x8),
            "testc_si128" => p!("sse41.ptestc", (i64x2, i64x2) -> i32),
            "testnzc_si128" => p!("sse41.ptestnzc", (i64x2, i64x2) -> i32),
            "testz_si128" => p!("sse41.ptestz", (i64x2, i64x2) -> i32),
            _ => return None
        })
    } else if name.starts_with("mm256_") {
        Some(match &name["mm256_".len()..] {
            _ => return None,
        })
    } else {
        None
    }
}
<commit_msg>Add most AVX intrinsics.<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use {Intrinsic, i, f, v};
use rustc::middle::ty;

macro_rules! p {
    ($name: expr, ($($inputs: tt),*) -> $output: tt) => {
        plain!(concat!("llvm.x86.", $name), ($($inputs),*) -> $output)
    }
}

pub fn find<'tcx>(_tcx: &ty::ctxt<'tcx>, name: &str) -> Option<Intrinsic> {
    if name.starts_with("mm_") {
        Some(match &name["mm_".len()..] {
            "sqrt_ps" => plain!("llvm.sqrt.v4f32", (f32x4) -> f32x4),
            "sqrt_pd" => plain!("llvm.sqrt.v2f64", (f64x2) -> f64x2),

            "movemask_ps" => p!("sse.movmsk.ps", (f32x4) -> i32),
            "max_ps" => p!("sse.max.ps", (f32x4, f32x4) -> f32x4),
            "min_ps" => p!("sse.min.ps", (f32x4, f32x4) -> f32x4),
            "rsqrt_ps" => p!("sse.rsqrt.ps", (f32x4) -> f32x4),
            "rcp_ps" => p!("sse.rcp.ps", (f32x4) -> f32x4),

            "adds_epi16" => p!("sse2.padds.w", (i16x8, i16x8) -> i16x8),
            "adds_epi8" => p!("sse2.padds.b", (i8x16, i8x16) -> i8x16),
            "adds_epu16" => p!("sse2.paddus.w", (i16x8, i16x8) -> i16x8),
            "adds_epu8" => p!("sse2.paddus.b", (i8x16, i8x16) -> i8x16),
            "avg_epu16" => p!("sse2.pavg.w", (i16x8, i16x8) -> i16x8),
            "avg_epu8" => p!("sse2.pavg.b", (i8x16, i8x16) -> i8x16),
            "madd_epi16" => p!("sse2.pmadd.wd", (i16x8, i16x8) -> i32x4),
            "max_epi16" => p!("sse2.pmaxs.w", (i16x8, i16x8) -> i16x8),
            "max_epu8" => p!("sse2.pmaxu.b", (i8x16, i8x16) -> i8x16),
            "max_pd" => p!("sse2.max.pd", (f64x2, f64x2) -> f64x2),
            "min_epi16" => p!("sse2.pmins.w", (i16x8, i16x8) -> i16x8),
            "min_epu8" => p!("sse2.pminu.b", (i8x16, i8x16) -> i8x16),
            "min_pd" => p!("sse2.min.pd", (f64x2, f64x2) -> f64x2),
            "movemask_pd" => p!("sse2.movmsk.pd", (f64x2) -> i32),
            "movemask_epi8" => p!("sse2.pmovmskb.128", (i8x16) -> i32),
            "mul_epu32" => p!("sse2.pmulu.dq", (i32x4, i32x4) -> i64x2),
            "mulhi_epi16" => p!("sse2.pmulh.w", (i8x16, i8x16) -> i8x16),
            "mulhi_epu16" => p!("sse2.pmulhu.w", (i8x16, i8x16) -> i8x16),
            "packs_epi16" => p!("sse2.packsswb.128", (i16x8, i16x8) -> i8x16),
            "packs_epi32" => p!("sse2.packssdw.128", (i32x4, i32x4) -> i16x8),
            "packus_epi16" => p!("sse2.packuswb.128", (i16x8, i16x8) -> i8x16),
            "sad_epu8" => p!("sse2.psad.bw", (i8x16, i8x16) -> i64x2),
            "subs_epi16" => p!("sse2.psubs.w", (i16x8, i16x8) -> i16x8),
            "subs_epi8" => p!("sse2.psubs.b", (i8x16, i8x16) -> i8x16),
            "subs_epu16" => p!("sse2.psubus.w", (i16x8, i16x8) -> i16x8),
            "subs_epu8" => p!("sse2.psubus.b", (i8x16, i8x16) -> i8x16),

            "addsub_pd" => p!("sse3.addsub.pd", (f64x2, f64x2) -> f64x2),
            "addsub_ps" => p!("sse3.addsub.ps", (f32x4, f32x4) -> f32x4),
            "hadd_pd" => p!("sse3.hadd.pd", (f64x2, f64x2) -> f64x2),
            "hadd_ps" => p!("sse3.hadd.ps", (f32x4, f32x4) -> f32x4),
            "hsub_pd" => p!("sse3.hsub.pd", (f64x2, f64x2) -> f64x2),
            "hsub_ps" => p!("sse3.hsub.ps", (f32x4, f32x4) -> f32x4),

            "abs_epi16" => p!("ssse3.pabs.w.128", (i16x8) -> i16x8),
            "abs_epi32" => p!("ssse3.pabs.d.128", (i32x4) -> i32x4),
            "abs_epi8" => p!("ssse3.pabs.b.128", (i8x16) -> i8x16),
            "hadd_epi16" => p!("ssse3.phadd.w.128", (i16x8, i16x8) -> i16x8),
            "hadd_epi32" => p!("ssse3.phadd.d.128", (i32x4, i32x4) -> i32x4),
            "hadds_epi16" => p!("ssse3.phadd.sw.128", (i16x8, i16x8) -> i16x8),
            "hsub_epi16" => p!("ssse3.phsub.w.128", (i16x8, i16x8) -> i16x8),
            "hsub_epi32" => p!("ssse3.phsub.d.128", (i32x4, i32x4) -> i32x4),
            "hsubs_epi16" => p!("ssse3.phsub.sw.128", (i16x8, i16x8) -> i16x8),
            "maddubs_epi16" => p!("ssse3.pmadd.ub.sw.128", (i8x16, i8x16) -> i16x8),
            "mulhrs_epi16" => p!("ssse3.pmul.hr.sw.128", (i16x8, i16x8) -> i16x8),
            "shuffle_epi8" => p!("ssse3.pshuf.b.128", (i8x16, i8x16) -> i8x16),
            "sign_epi16" => p!("ssse3.psign.w.128", (i16x8, i16x8) -> i16x8),
            "sign_epi32" => p!("ssse3.psign.d.128", (i32x4, i32x4) -> i32x4),
            "sign_epi8" => p!("ssse3.psign.b.128", (i8x16, i8x16) -> i8x16),

            "max_epi32" => p!("sse41.pmaxsd", (i32x4, i32x4) -> i32x4),
            "max_epi8" => p!("sse41.pmaxsb", (i8x16, i8x16) -> i8x16),
            "max_epu16" => p!("sse41.pmaxuw", (i16x8, i16x8) -> i16x8),
            "max_epu32" => p!("sse41.pmaxud", (i32x4, i32x4) -> i32x4),
            "min_epi32" => p!("sse41.pminsd", (i32x4, i32x4) -> i32x4),
            "min_epi8" => p!("sse41.pminsb", (i8x16, i8x16) -> i8x16),
            "min_epu16" => p!("sse41.pminuw", (i16x8, i16x8) -> i16x8),
            "min_epu32" => p!("sse41.pminud", (i32x4, i32x4) -> i32x4),
            "minpos_epu16" => p!("sse41.phminposuw", (i16x8) -> i16x8),
            "mul_epi32" => p!("sse41.muldq", (i32x4, i32x4) -> i64x2),
            "packus_epi32" => p!("sse41.packusdw", (i32x4, i32x4) -> i16x8),
            "testc_si128" => p!("sse41.ptestc", (i64x2, i64x2) -> i32),
            "testnzc_si128" => p!("sse41.ptestnzc", (i64x2, i64x2) -> i32),
            "testz_si128" => p!("sse41.ptestz", (i64x2, i64x2) -> i32),

            "permutevar_pd" => p!("avx.vpermilvar.pd", (f64x2, i64x2) -> f64x2),
            "permutevar_ps" => p!("avx.vpermilvar.ps", (f32x4, i32x4) -> f32x4),
            "testc_pd" => p!("avx.vtestc.pd", (f64x2, f64x2) -> i32),
            "testc_ps" => p!("avx.vtestc.ps", (f32x4, f32x4) -> i32),
            "testnzc_pd" => p!("avx.vtestnzc.pd", (f64x2, f64x2) -> i32),
            "testnzc_ps" => p!("avx.vtestnzc.ps", (f32x4, f32x4) -> i32),
            "testz_pd" => p!("avx.vtestz.pd", (f64x2, f64x2) -> i32),
            "testz_ps" => p!("avx.vtestz.ps", (f32x4, f32x4) -> i32),

            _ => return None
        })
    } else if name.starts_with("mm256_") {
        Some(match &name["mm256_".len()..] {
            "addsub_pd" => p!("avx.addsub.pd.256", (f64x4, f64x4) -> f64x4),
            "addsub_ps" => p!("avx.addsub.ps.256", (f32x8, f32x8) -> f32x8),
            "hadd_pd" => p!("avx.hadd.pd.256", (f64x4, f64x4) -> f64x4),
            "hadd_ps" => p!("avx.hadd.ps.256", (f32x8, f32x8) -> f32x8),
            "hsub_pd" => p!("avx.hsub.pd.256", (f64x4, f64x4) -> f64x4),
            "hsub_ps" => p!("avx.hsub.ps.256", (f32x8, f32x8) -> f32x8),
            "max_pd" => p!("avx.max.pd.256", (f64x4, f64x4) -> f64x4),
            "max_ps" => p!("avx.max.ps.256", (f32x8, f32x8) -> f32x8),
            "min_pd" => p!("avx.min.pd.256", (f64x4, f64x4) -> f64x4),
            "min_ps" => p!("avx.min.ps.256", (f32x8, f32x8) -> f32x8),
            "permutevar_pd" => p!("avx.vpermilvar.pd.256", (f64x4, i64x4) -> f64x4),
            "permutevar_ps" => p!("avx.vpermilvar.ps.256", (f32x8, i32x8) -> f32x8),
            "rcp_ps" => p!("avx.rcp.ps.256", (f32x8) -> f32x8),
            "rsqrt_ps" => p!("avx.rsqrt.ps.256", (f32x8) -> f32x8),
            "sqrt_pd" => p!("llvm.sqrt.v4f64", (f64x4) -> f64x4),
            "sqrt_ps" => p!("llvm.sqrt.v8f32", (f32x8) -> f32x8),
            "testc_pd" => p!("avx.vtestc.pd.256", (f64x4, f64x4) -> i32),
            "testc_ps" => p!("avx.vtestc.ps.256", (f32x8, f32x8) -> i32),
            "testnzc_pd" => p!("avx.vtestnzc.pd.256", (f64x4, f64x4) -> i32),
            "testnzc_ps" => p!("avx.vtestnzc.ps.256", (f32x8, f32x8) -> i32),
            "testz_pd" => p!("avx.vtestz.pd.256", (f64x4, f64x4) -> i32),
            "testz_ps" => p!("avx.vtestz.ps.256", (f32x8, f32x8) -> i32),
            _ => return None,
        })
    } else {
        None
    }
}
<|endoftext|>


<commit_before>extern crate zim;
extern crate clap;
extern crate stopwatch;
extern crate crossbeam;

use std::fs::File;
use std::io::{Write, BufWriter};
use std::error::Error;
use std::collections::HashMap;
use std::path::Path;

use clap::{Arg, App};
use stopwatch::Stopwatch;

use zim::{Zim, Target};

fn main() {
    let matches = App::new("zimextractor")
        .version("0.1")
        .about("Extract zim files")
        .arg(Arg::with_name("out")
                 .short("o")
                 .long("out")
                 .help("Output directory")
                 .takes_value(true))
        .arg(Arg::with_name("skip-link")
                 .long("skip-link")
                 .help("Skip genrating hard links")
                 .takes_value(false))
        .arg(Arg::with_name("INPUT")
                 .help("Set the zim file to extract")
                 .required(true)
                 .index(1))
        .get_matches();

    let skip_link = matches.is_present("skip-link");
    let out = matches.value_of("out").unwrap_or("out");
    let root_output = Path::new(out);

    let input = matches.value_of("INPUT").unwrap();

    println!("Extracting file: {} to {}\n", input, out);

    let sw = Stopwatch::start_new();

    let zim = Zim::new(input).ok().unwrap();

    // map between cluster and directory entry
    let mut cluster_map = HashMap::new();

    println!("  Creating cluster map");
    for i in zim.iterate_by_urls() {
        if let Some(Target::Cluster(cid, _)) = i.target {
            cluster_map.entry(cid).or_insert(Vec::new()).push(i);
        }
    }

    println!("  Extracting entries: {}", zim.header.cluster_count);

    crossbeam::scope(|scope| for (cid, entries) in cluster_map {
                         let mut cluster = zim.get_cluster(cid).unwrap();

                         scope.spawn(move || {
            cluster.decompress();

            for entry in entries {
                if let Some(Target::Cluster(_cid, bid)) = entry.target {
                    assert_eq!(cid, _cid);
                    let mut s = String::new();
                    s.push(entry.namespace);
                    let out_path = root_output.join(&s).join(&entry.url);
                    safe_write(&out_path, cluster.get_blob(bid), 0);
                }
            }
        });
                     });

    println!("  Extraction done in {}ms", sw.elapsed_ms());

    if !skip_link {
        println!("  Linking redirects");

        // link all redirects
        for entry in zim.iterate_by_urls() {
            // get redirect entry
            if let Some(Target::Redirect(redir)) = entry.target {
                let redir = zim.get_by_url_index(redir).unwrap();

                let mut s = String::new();
                s.push(redir.namespace);
                let src = root_output.join(&s).join(&redir.url);

                let mut d = String::new();
                d.push(entry.namespace);
                let dst = root_output.join(&s).join(&entry.url);

                if !dst.exists() {
                    // println!("{:?} -> {:?}", src, dst);
                    std::fs::hard_link(src, dst).unwrap();
                }
            }
        }

        println!("  Linking done in {}ms", sw.elapsed_ms());
    } else {
        println!("  Skipping linking...");
    }

    if let Some(main_page_idx) = zim.header.main_page {
        let page = zim.get_by_url_index(main_page_idx).unwrap();
        println!("  Main page is {}", page.url);
    }
}

fn safe_write(path: &Path, data: &[u8], count: usize) {
    let display = path.display();
    let contain_path = path.parent().unwrap();

    ensure_dir(contain_path);

    match File::create(&path) {
        Err(why) => {
            println!("couldn't create {}: {}", display, why.description());

            if count == 1 {
                safe_write(path, data, 2);
            } else {
                panic!("failed retry: couldn't create {}: {}",
                       display,
                       why.description());
            }
        }
        Ok(file) => {
            let mut writer = BufWriter::new(&file);

            if let Err(why) = writer.write_all(data) {
                println!("couldn't write to {}: {}", display, why.description());
            }
        }
    }
}

fn ensure_dir(path: &Path) {
    match std::fs::create_dir_all(path) {
        Err(e) => {
            use std::io::ErrorKind::*;

            match e.kind() {
                // do not panic if it already exists, that's fine, we just want to make
                // sure we have it before moving on
                AlreadyExists => {}
                _ => {
                    panic!("failed to create {}: {}", path.display(), e.description());
                }
            }
        }
        _ => {}
    }
}
<commit_msg>fix retry counting<commit_after>extern crate zim;
extern crate clap;
extern crate stopwatch;
extern crate crossbeam;

use std::fs::File;
use std::io::{Write, BufWriter};
use std::error::Error;
use std::collections::HashMap;
use std::path::Path;

use clap::{Arg, App};
use stopwatch::Stopwatch;

use zim::{Zim, Target};

fn main() {
    let matches = App::new("zimextractor")
        .version("0.1")
        .about("Extract zim files")
        .arg(Arg::with_name("out")
                 .short("o")
                 .long("out")
                 .help("Output directory")
                 .takes_value(true))
        .arg(Arg::with_name("skip-link")
                 .long("skip-link")
                 .help("Skip genrating hard links")
                 .takes_value(false))
        .arg(Arg::with_name("INPUT")
                 .help("Set the zim file to extract")
                 .required(true)
                 .index(1))
        .get_matches();

    let skip_link = matches.is_present("skip-link");
    let out = matches.value_of("out").unwrap_or("out");
    let root_output = Path::new(out);

    let input = matches.value_of("INPUT").unwrap();

    println!("Extracting file: {} to {}\n", input, out);

    let sw = Stopwatch::start_new();

    let zim = Zim::new(input).ok().unwrap();

    // map between cluster and directory entry
    let mut cluster_map = HashMap::new();

    println!("  Creating cluster map");
    for i in zim.iterate_by_urls() {
        if let Some(Target::Cluster(cid, _)) = i.target {
            cluster_map.entry(cid).or_insert(Vec::new()).push(i);
        }
    }

    println!("  Extracting entries: {}", zim.header.cluster_count);

    crossbeam::scope(|scope| for (cid, entries) in cluster_map {
                         let mut cluster = zim.get_cluster(cid).unwrap();

                         scope.spawn(move || {
            cluster.decompress();

            for entry in entries {
                if let Some(Target::Cluster(_cid, bid)) = entry.target {
                    assert_eq!(cid, _cid);
                    let mut s = String::new();
                    s.push(entry.namespace);
                    let out_path = root_output.join(&s).join(&entry.url);
                    safe_write(&out_path, cluster.get_blob(bid), 1);
                }
            }
        });
                     });

    println!("  Extraction done in {}ms", sw.elapsed_ms());

    if !skip_link {
        println!("  Linking redirects");

        // link all redirects
        for entry in zim.iterate_by_urls() {
            // get redirect entry
            if let Some(Target::Redirect(redir)) = entry.target {
                let redir = zim.get_by_url_index(redir).unwrap();

                let mut s = String::new();
                s.push(redir.namespace);
                let src = root_output.join(&s).join(&redir.url);

                let mut d = String::new();
                d.push(entry.namespace);
                let dst = root_output.join(&s).join(&entry.url);

                if !dst.exists() {
                    // println!("{:?} -> {:?}", src, dst);
                    std::fs::hard_link(src, dst).unwrap();
                }
            }
        }

        println!("  Linking done in {}ms", sw.elapsed_ms());
    } else {
        println!("  Skipping linking...");
    }

    if let Some(main_page_idx) = zim.header.main_page {
        let page = zim.get_by_url_index(main_page_idx).unwrap();
        println!("  Main page is {}", page.url);
    }
}

fn safe_write(path: &Path, data: &[u8], count: usize) {
    let display = path.display();
    let contain_path = path.parent().unwrap();

    ensure_dir(contain_path);

    match File::create(&path) {
        Err(why) => {
            println!("couldn't create {}: {}", display, why.description());

            if count < 3 {
                safe_write(path, data, count + 1);
            } else {
                panic!("failed retry: couldn't create {}: {}",
                       display,
                       why.description());
            }
        }
        Ok(file) => {
            let mut writer = BufWriter::new(&file);

            if let Err(why) = writer.write_all(data) {
                println!("couldn't write to {}: {}", display, why.description());
            }
        }
    }
}

fn ensure_dir(path: &Path) {
    match std::fs::create_dir_all(path) {
        Err(e) => {
            use std::io::ErrorKind::*;

            match e.kind() {
                // do not panic if it already exists, that's fine, we just want to make
                // sure we have it before moving on
                AlreadyExists => {}
                _ => {
                    panic!("failed to create {}: {}", path.display(), e.description());
                }
            }
        }
        _ => {}
    }
}
<|endoftext|>


<commit_before>use crate::custom::credential::NewAwsCredsForStsCreds;
use crate::{AssumeRoleWithWebIdentityRequest, Sts, StsClient};
use rusoto_core::credential::{
    AwsCredentials, CredentialsError, ProvideAwsCredentials, Secret, Variable,
};
use rusoto_core::request::HttpClient;
use rusoto_core::{Client, Region};

use async_trait::async_trait;

const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";

const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";

const AWS_ROLE_SESSION_NAME: &str = "AWS_ROLE_SESSION_NAME";

/// WebIdentityProvider using OpenID Connect bearer token to retrieve AWS IAM credentials.
///
/// See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html for
/// more details.
#[derive(Debug, Clone)]
pub struct WebIdentityProvider {
    /// The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// Your application must get this token by authenticating the user who is using your application
    /// with a web identity provider before the application makes an AssumeRoleWithWebIdentity call.
    pub web_identity_token: Variable<Secret, CredentialsError>,
    /// The Amazon Resource Name (ARN) of the role that the caller is assuming.
    pub role_arn: Variable<String, CredentialsError>,
    /// An identifier for the assumed role session. Typically, you pass the name or identifier that is
    /// associated with the user who is using your application. That way, the temporary security credentials
    /// that your application will use are associated with that user. This session name is included as part
    /// of the ARN and assumed role ID in the AssumedRoleUser response element.
    pub role_session_name: Variable<String, CredentialsError>,
}

impl WebIdentityProvider {
    /// Create new WebIdentityProvider by explicitly passing its configuration.
    pub fn new<A, B, C>(web_identity_token: A, role_arn: B, role_session_name: Option<C>) -> Self
    where
        A: Into<Variable<Secret, CredentialsError>>,
        B: Into<Variable<String, CredentialsError>>,
        C: Into<Variable<String, CredentialsError>>,
    {
        Self {
            web_identity_token: web_identity_token.into(),
            role_arn: role_arn.into(),
            role_session_name: role_session_name
                .map(|v| v.into())
                .unwrap_or_else(|| Variable::with_value(Self::create_session_name())),
        }
    }

    /// Creat a WebIdentityProvider from the following environment variables:
    ///
    /// - `AWS_WEB_IDENTITY_TOKEN_FILE` path to the web identity token file.
    /// - `AWS_ROLE_ARN` ARN of the role to assume.
    /// - `AWS_ROLE_SESSION_NAME` (optional) name applied to the assume-role session.
    ///
    /// See https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts-technical-overview.html
    /// for more information about how IAM Roles for Kubernetes Service Accounts works.
    pub fn from_k8s_env() -> Self {
        Self::_from_k8s_env(
            Variable::from_env_var(AWS_WEB_IDENTITY_TOKEN_FILE),
            Variable::from_env_var(AWS_ROLE_ARN),
            Some(Variable::from_env_var(AWS_ROLE_SESSION_NAME)),
        )
    }

    /// Used by unit testing
    pub(crate) fn _from_k8s_env(
        token_file: Variable<String, CredentialsError>,
        role: Variable<String, CredentialsError>,
        session_name: Option<Variable<String, CredentialsError>>,
    ) -> Self {
        Self::new(
            Variable::dynamic(move || Variable::from_text_file(token_file.resolve()?).resolve()),
            role,
            session_name,
        )
    }

    #[cfg(test)]
    pub(crate) fn load_token(&self) -> Result<Secret, CredentialsError> {
        self.web_identity_token.resolve()
    }

    fn create_session_name() -> String {
        // TODO can we do better here?
        // - Pod service account, Pod name and Pod namespace
        // - EC2 Instance ID if available
        // - IP address if available
        // - ...
        // Having some information in the session name that identifies the client would enable
        // better correlation analysis in CloudTrail.
        "WebIdentitySession".to_string()
    }
}

#[async_trait]
impl ProvideAwsCredentials for WebIdentityProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let http_client = match HttpClient::new() {
            Ok(c) => c,
            Err(e) => return Err(CredentialsError::new(e)),
        };
        let client = Client::new_not_signing(http_client);
        let sts = StsClient::new_with_client(client, Region::default());
        let mut req = AssumeRoleWithWebIdentityRequest::default();

        req.role_arn = self.role_arn.resolve()?;
        req.web_identity_token = self.web_identity_token.resolve()?.to_string();
        req.role_session_name = self.role_session_name.resolve()?;

        let assume_role = sts.assume_role_with_web_identity(req).await;
        match assume_role {
            Err(e) => Err(CredentialsError::new(e)),
            Ok(role) => match role.credentials {
                None => Err(CredentialsError::new(format!(
                    "No credentials found in AssumeRoleWithWebIdentityResponse: {:?}",
                    role
                ))),
                Some(c) => AwsCredentials::new_for_credentials(c),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn api_ergonomy() {
        WebIdentityProvider::new(Secret::from("".to_string()), "", Some("".to_string()));
    }

    #[test]
    fn from_k8s_env() -> Result<(), CredentialsError> {
        const TOKEN_VALUE: &str = "secret";
        const ROLE_ARN: &str = "role";
        const SESSION_NAME: &str = "session";
        let mut file = NamedTempFile::new()?;
        // We use writeln to add an extra newline at the end of the token, which should be
        // removed by Variable::from_text_file.
        writeln!(file, "{}", TOKEN_VALUE)?;
        let p = WebIdentityProvider::_from_k8s_env(
            Variable::with_value(file.path().to_string_lossy().to_string()),
            Variable::with_value(ROLE_ARN.to_string()),
            Some(Variable::with_value(SESSION_NAME.to_string())),
        );
        let token = p.load_token()?;
        assert_eq!(token.as_ref(), TOKEN_VALUE);
        Ok(())
    }
}
<commit_msg>Rollup merge of #1694<commit_after>use crate::custom::credential::NewAwsCredsForStsCreds;
use crate::{AssumeRoleWithWebIdentityRequest, Sts, StsClient};
use rusoto_core::credential::{
    AwsCredentials, CredentialsError, ProvideAwsCredentials, Secret, Variable,
};
use rusoto_core::request::HttpClient;
use rusoto_core::{Client, Region};

use async_trait::async_trait;

const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";

const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";

const AWS_ROLE_SESSION_NAME: &str = "AWS_ROLE_SESSION_NAME";

/// WebIdentityProvider using OpenID Connect bearer token to retrieve AWS IAM credentials.
///
/// See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html for
/// more details.
#[derive(Debug, Clone)]
pub struct WebIdentityProvider {
    /// The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// Your application must get this token by authenticating the user who is using your application
    /// with a web identity provider before the application makes an AssumeRoleWithWebIdentity call.
    pub web_identity_token: Variable<Secret, CredentialsError>,
    /// The Amazon Resource Name (ARN) of the role that the caller is assuming.
    pub role_arn: Variable<String, CredentialsError>,
    /// An identifier for the assumed role session. Typically, you pass the name or identifier that is
    /// associated with the user who is using your application. That way, the temporary security credentials
    /// that your application will use are associated with that user. This session name is included as part
    /// of the ARN and assumed role ID in the AssumedRoleUser response element.
    pub role_session_name: Variable<String, CredentialsError>,
}

impl WebIdentityProvider {
    /// Create new WebIdentityProvider by explicitly passing its configuration.
    pub fn new<A, B, C>(web_identity_token: A, role_arn: B, role_session_name: Option<C>) -> Self
    where
        A: Into<Variable<Secret, CredentialsError>>,
        B: Into<Variable<String, CredentialsError>>,
        C: Into<Variable<String, CredentialsError>>,
    {
        Self {
            web_identity_token: web_identity_token.into(),
            role_arn: role_arn.into(),
            role_session_name: role_session_name
                .map(|v| v.into())
                .unwrap_or_else(|| Variable::with_value(Self::create_session_name())),
        }
    }

    /// Creat a WebIdentityProvider from the following environment variables:
    ///
    /// - `AWS_WEB_IDENTITY_TOKEN_FILE` path to the web identity token file.
    /// - `AWS_ROLE_ARN` ARN of the role to assume.
    /// - `AWS_ROLE_SESSION_NAME` (optional) name applied to the assume-role session.
    ///
    /// See https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts-technical-overview.html
    /// for more information about how IAM Roles for Kubernetes Service Accounts works.
    pub fn from_k8s_env() -> Self {
        Self::_from_k8s_env(
            Variable::from_env_var(AWS_WEB_IDENTITY_TOKEN_FILE),
            Variable::from_env_var(AWS_ROLE_ARN),
            Some(Variable::from_env_var(AWS_ROLE_SESSION_NAME)),
        )
    }

    /// Used by unit testing
    pub(crate) fn _from_k8s_env(
        token_file: Variable<String, CredentialsError>,
        role: Variable<String, CredentialsError>,
        session_name: Option<Variable<String, CredentialsError>>,
    ) -> Self {
        Self::new(
            Variable::dynamic(move || Variable::from_text_file(token_file.resolve()?).resolve()),
            role,
            session_name,
        )
    }

    #[cfg(test)]
    pub(crate) fn load_token(&self) -> Result<Secret, CredentialsError> {
        self.web_identity_token.resolve()
    }

    fn create_session_name() -> String {
        // TODO can we do better here?
        // - Pod service account, Pod name and Pod namespace
        // - EC2 Instance ID if available
        // - IP address if available
        // - ...
        // Having some information in the session name that identifies the client would enable
        // better correlation analysis in CloudTrail.
        "WebIdentitySession".to_string()
    }
}

#[async_trait]
impl ProvideAwsCredentials for WebIdentityProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let http_client = match HttpClient::new() {
            Ok(c) => c,
            Err(e) => return Err(CredentialsError::new(e)),
        };
        let client = Client::new_not_signing(http_client);
        let sts = StsClient::new_with_client(client, Region::default());
        let mut req = AssumeRoleWithWebIdentityRequest::default();

        req.role_arn = self.role_arn.resolve()?;
        req.web_identity_token = self.web_identity_token.resolve()?.as_ref().to_string();
        req.role_session_name = self.role_session_name.resolve()?;

        let assume_role = sts.assume_role_with_web_identity(req).await;
        match assume_role {
            Err(e) => Err(CredentialsError::new(e)),
            Ok(role) => match role.credentials {
                None => Err(CredentialsError::new(format!(
                    "No credentials found in AssumeRoleWithWebIdentityResponse: {:?}",
                    role
                ))),
                Some(c) => AwsCredentials::new_for_credentials(c),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn api_ergonomy() {
        WebIdentityProvider::new(Secret::from("".to_string()), "", Some("".to_string()));
    }

    #[test]
    fn from_k8s_env() -> Result<(), CredentialsError> {
        const TOKEN_VALUE: &str = "secret";
        const ROLE_ARN: &str = "role";
        const SESSION_NAME: &str = "session";
        let mut file = NamedTempFile::new()?;
        // We use writeln to add an extra newline at the end of the token, which should be
        // removed by Variable::from_text_file.
        writeln!(file, "{}", TOKEN_VALUE)?;
        let p = WebIdentityProvider::_from_k8s_env(
            Variable::with_value(file.path().to_string_lossy().to_string()),
            Variable::with_value(ROLE_ARN.to_string()),
            Some(Variable::with_value(SESSION_NAME.to_string())),
        );
        let token = p.load_token()?;
        assert_eq!(token.as_ref(), TOKEN_VALUE);
        Ok(())
    }
}
<|endoftext|>


<commit_before>use either::Either;
use ide_db::syntax_helpers::node_ext::walk_ty;
use itertools::Itertools;
use syntax::{
    ast::{self, edit::IndentLevel, AstNode, HasGenericParams, HasName},
    match_ast,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    if ctx.has_empty_selection() {
        return None;
    }

    let ty = ctx.find_node_at_range::<ast::Type>()?;
    let item = ty.syntax().ancestors().find_map(ast::Item::cast)?;
    let assoc_owner = item.syntax().ancestors().nth(2).and_then(|it| {
        match_ast! {
            match it {
                ast::Trait(tr) => Some(Either::Left(tr)),
                ast::Impl(impl_) => Some(Either::Right(impl_)),
                _ => None,
            }
        }
    });
    let node = assoc_owner.as_ref().map_or_else(
        || item.syntax(),
        |impl_| impl_.as_ref().either(AstNode::syntax, AstNode::syntax),
    );
    let insert_pos = node.text_range().start();
    let target = ty.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            let mut known_generics = match item.generic_param_list() {
                Some(it) => it.generic_params().collect(),
                None => Vec::new(),
            };
            if let Some(it) = assoc_owner.as_ref().and_then(|it| match it {
                Either::Left(it) => it.generic_param_list(),
                Either::Right(it) => it.generic_param_list(),
            }) {
                known_generics.extend(it.generic_params());
            }
            let generics = collect_used_generics(&ty, &known_generics);

            let replacement = if !generics.is_empty() {
                format!(
                    "Type<{}>",
                    generics.iter().format_with(", ", |generic, f| {
                        match generic {
                            ast::GenericParam::ConstParam(cp) => f(&cp.name().unwrap()),
                            ast::GenericParam::LifetimeParam(lp) => f(&lp.lifetime().unwrap()),
                            ast::GenericParam::TypeParam(tp) => f(&tp.name().unwrap()),
                        }
                    })
                )
            } else {
                String::from("Type")
            };
            builder.replace(target, replacement);

            let indent = IndentLevel::from_node(node);
            let generics = if !generics.is_empty() {
                format!("<{}>", generics.iter().format(", "))
            } else {
                String::new()
            };
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("type $0Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
                None => {
                    builder.insert(
                        insert_pos,
                        format!("type Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
            }
        },
    )
}

fn collect_used_generics<'gp>(
    ty: &ast::Type,
    known_generics: &'gp [ast::GenericParam],
) -> Vec<&'gp ast::GenericParam> {
    // can't use a closure -> closure here cause lifetime inference fails for that
    fn find_lifetime(text: &str) -> impl Fn(&&ast::GenericParam) -> bool + '_ {
        move |gp: &&ast::GenericParam| match gp {
            ast::GenericParam::LifetimeParam(lp) => {
                lp.lifetime().map_or(false, |lt| lt.text() == text)
            }
            _ => false,
        }
    }

    let mut generics = Vec::new();
    walk_ty(ty, &mut |ty| match ty {
        ast::Type::PathType(ty) => {
            if let Some(path) = ty.path() {
                if let Some(name_ref) = path.as_single_name_ref() {
                    if let Some(param) = known_generics.iter().find(|gp| {
                        match gp {
                            ast::GenericParam::ConstParam(cp) => cp.name(),
                            ast::GenericParam::TypeParam(tp) => tp.name(),
                            _ => None,
                        }
                        .map_or(false, |n| n.text() == name_ref.text())
                    }) {
                        generics.push(param);
                    }
                }
                generics.extend(
                    path.segments()
                        .filter_map(|seg| seg.generic_arg_list())
                        .flat_map(|it| it.generic_args())
                        .filter_map(|it| match it {
                            ast::GenericArg::LifetimeArg(lt) => {
                                let lt = lt.lifetime()?;
                                known_generics.iter().find(find_lifetime(&lt.text()))
                            }
                            _ => None,
                        }),
                );
            }
        }
        ast::Type::ImplTraitType(impl_ty) => {
            if let Some(it) = impl_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::DynTraitType(dyn_ty) => {
            if let Some(it) = dyn_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::RefType(ref_) => generics.extend(
            ref_.lifetime().and_then(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
        ),
        _ => (),
    });
    // stable resort to lifetime, type, const
    generics.sort_by_key(|gp| match gp {
        ast::GenericParam::ConstParam(_) => 2,
        ast::GenericParam::LifetimeParam(_) => 0,
        ast::GenericParam::TypeParam(_) => 1,
    });
    generics
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }

    #[test]
    fn extract_from_impl_or_trait() {
        // When invoked in an impl/trait, extracted type alias should be placed next to the
        // impl/trait, not inside.
        check_assist(
            extract_type_alias,
            r#"
impl S {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

impl S {
    fn f() -> Type {}
}
            "#,
        );
        check_assist(
            extract_type_alias,
            r#"
trait Tr {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

trait Tr {
    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn indentation() {
        check_assist(
            extract_type_alias,
            r#"
mod m {
    fn f() -> $0u8$0 {}
}
            "#,
            r#"
mod m {
    type $0Type = u8;

    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn generics() {
        check_assist(
            extract_type_alias,
            r#"
struct Struct<const C: usize>;
impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: $0&(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ())$0) {}
}
"#,
            r#"
struct Struct<const C: usize>;
type $0Type<'inner, 'outer, Outer, Inner, const INNER: usize, const OUTER: usize> = &(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ());

impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: Type<'inner, 'outer, Outer, Inner, INNER, OUTER>) {}
}
"#,
        );
    }
}
<commit_msg>Auto merge of #13110 - DesmondWillowbrook:issue-11197, r=jonas-schievink<commit_after>use either::Either;
use ide_db::syntax_helpers::node_ext::walk_ty;
use itertools::Itertools;
use syntax::{
    ast::{self, edit::IndentLevel, AstNode, HasGenericParams, HasName},
    match_ast,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    if ctx.has_empty_selection() {
        return None;
    }

    let ty = ctx.find_node_at_range::<ast::Type>()?;
    let item = ty.syntax().ancestors().find_map(ast::Item::cast)?;
    let assoc_owner = item.syntax().ancestors().nth(2).and_then(|it| {
        match_ast! {
            match it {
                ast::Trait(tr) => Some(Either::Left(tr)),
                ast::Impl(impl_) => Some(Either::Right(impl_)),
                _ => None,
            }
        }
    });
    let node = assoc_owner.as_ref().map_or_else(
        || item.syntax(),
        |impl_| impl_.as_ref().either(AstNode::syntax, AstNode::syntax),
    );
    let insert_pos = node.text_range().start();
    let target = ty.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            let mut known_generics = match item.generic_param_list() {
                Some(it) => it.generic_params().collect(),
                None => Vec::new(),
            };
            if let Some(it) = assoc_owner.as_ref().and_then(|it| match it {
                Either::Left(it) => it.generic_param_list(),
                Either::Right(it) => it.generic_param_list(),
            }) {
                known_generics.extend(it.generic_params());
            }
            let generics = collect_used_generics(&ty, &known_generics);

            let replacement = if !generics.is_empty() {
                format!(
                    "Type<{}>",
                    generics.iter().format_with(", ", |generic, f| {
                        match generic {
                            ast::GenericParam::ConstParam(cp) => f(&cp.name().unwrap()),
                            ast::GenericParam::LifetimeParam(lp) => f(&lp.lifetime().unwrap()),
                            ast::GenericParam::TypeParam(tp) => f(&tp.name().unwrap()),
                        }
                    })
                )
            } else {
                String::from("Type")
            };
            builder.replace(target, replacement);

            let indent = IndentLevel::from_node(node);
            let generics = if !generics.is_empty() {
                format!("<{}>", generics.iter().format(", "))
            } else {
                String::new()
            };
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("type $0Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
                None => {
                    builder.insert(
                        insert_pos,
                        format!("type Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
            }
        },
    )
}

fn collect_used_generics<'gp>(
    ty: &ast::Type,
    known_generics: &'gp [ast::GenericParam],
) -> Vec<&'gp ast::GenericParam> {
    // can't use a closure -> closure here cause lifetime inference fails for that
    fn find_lifetime(text: &str) -> impl Fn(&&ast::GenericParam) -> bool + '_ {
        move |gp: &&ast::GenericParam| match gp {
            ast::GenericParam::LifetimeParam(lp) => {
                lp.lifetime().map_or(false, |lt| lt.text() == text)
            }
            _ => false,
        }
    }

    let mut generics = Vec::new();
    walk_ty(ty, &mut |ty| match ty {
        ast::Type::PathType(ty) => {
            if let Some(path) = ty.path() {
                if let Some(name_ref) = path.as_single_name_ref() {
                    if let Some(param) = known_generics.iter().find(|gp| {
                        match gp {
                            ast::GenericParam::ConstParam(cp) => cp.name(),
                            ast::GenericParam::TypeParam(tp) => tp.name(),
                            _ => None,
                        }
                        .map_or(false, |n| n.text() == name_ref.text())
                    }) {
                        generics.push(param);
                    }
                }
                generics.extend(
                    path.segments()
                        .filter_map(|seg| seg.generic_arg_list())
                        .flat_map(|it| it.generic_args())
                        .filter_map(|it| match it {
                            ast::GenericArg::LifetimeArg(lt) => {
                                let lt = lt.lifetime()?;
                                known_generics.iter().find(find_lifetime(&lt.text()))
                            }
                            _ => None,
                        }),
                );
            }
        }
        ast::Type::ImplTraitType(impl_ty) => {
            if let Some(it) = impl_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::DynTraitType(dyn_ty) => {
            if let Some(it) = dyn_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::RefType(ref_) => generics.extend(
            ref_.lifetime().and_then(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
        ),
        ast::Type::ArrayType(ar) => {
            if let Some(expr) = ar.expr() {
                if let ast::Expr::PathExpr(p) = expr {
                    if let Some(path) = p.path() {
                        if let Some(name_ref) = path.as_single_name_ref() {
                            if let Some(param) = known_generics.iter().find(|gp| {
                                if let ast::GenericParam::ConstParam(cp) = gp {
                                    cp.name().map_or(false, |n| n.text() == name_ref.text())
                                } else {
                                    false
                                }
                            }) {
                                generics.push(param);
                            }
                        }
                    }
                }
            }
        }
        _ => (),
    });
    // stable resort to lifetime, type, const
    generics.sort_by_key(|gp| match gp {
        ast::GenericParam::ConstParam(_) => 2,
        ast::GenericParam::LifetimeParam(_) => 0,
        ast::GenericParam::TypeParam(_) => 1,
    });
    generics
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }

    #[test]
    fn extract_from_impl_or_trait() {
        // When invoked in an impl/trait, extracted type alias should be placed next to the
        // impl/trait, not inside.
        check_assist(
            extract_type_alias,
            r#"
impl S {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

impl S {
    fn f() -> Type {}
}
            "#,
        );
        check_assist(
            extract_type_alias,
            r#"
trait Tr {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

trait Tr {
    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn indentation() {
        check_assist(
            extract_type_alias,
            r#"
mod m {
    fn f() -> $0u8$0 {}
}
            "#,
            r#"
mod m {
    type $0Type = u8;

    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn generics() {
        check_assist(
            extract_type_alias,
            r#"
struct Struct<const C: usize>;
impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: $0&(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ())$0) {}
}
"#,
            r#"
struct Struct<const C: usize>;
type $0Type<'inner, 'outer, Outer, Inner, const INNER: usize, const OUTER: usize> = &(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ());

impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: Type<'inner, 'outer, Outer, Inner, INNER, OUTER>) {}
}
"#,
        );
    }

    #[test]
    fn issue_11197() {
        check_assist(
            extract_type_alias,
            r#"
struct Foo<T, const N: usize>
where
    [T; N]: Sized,
{
    arr: $0[T; N]$0,
}
            "#,
            r#"
type $0Type<T, const N: usize> = [T; N];

struct Foo<T, const N: usize>
where
    [T; N]: Sized,
{
    arr: Type<T, N>,
}
            "#,
        );
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2104 Michael Dvorkin
// Ante is an esoteric programming language where all you've got is a deck of cards.
//
// This is Ante implementation in Rust.

extern crate regex;
extern crate num;

use std::io::File;
use std::collections::HashMap;
use std::num::Zero;
use std::num::from_int;
use std::str::{is_utf8, from_utf8};
use regex::Regex;
use num::bigint::BigInt;

const A: uint = 'A' as uint;
const J: uint = 'J' as uint;
const Q: uint = 'Q' as uint;
const K: uint = 'K' as uint;
const D: uint = '♦' as uint;
const H: uint = '♥' as uint;
const S: uint = '♠' as uint;
const C: uint = '♣' as uint;

struct Card {
    rank: uint,
    suit: uint
}

struct Ante {
    pc:     uint,                   // Program counter (index within ante.code)
    line:   uint,                   // Current line number.
    code:   Vec<Card>,              // Array of cards.
    vars:   HashMap<uint, BigInt>,  // Four registers hashed by suit.
    labels: HashMap<uint, uint>,    // Labels for ante.pc to jump to.
    buffer: Vec<u8>                 // Buffer to collect UTF-8 character bytes.
}

impl Ante {
    fn new(filename: &str) -> Ante {
        let mut vars: HashMap<uint, BigInt> = HashMap::new();
        vars.insert(D, Ante::big(0));
        vars.insert(H, Ante::big(0));
        vars.insert(S, Ante::big(0));
        vars.insert(C, Ante::big(0));

        let code = Ante::parse(filename);
        let labels = Ante::resolve(&code);

        Ante {
            pc:     0,
            line:   0,
            code:   code,
            vars:   vars,
            labels: labels,
            buffer: vec![]
        }
    }

    fn run(&mut self) -> &Ante {
        while self.pc < self.code.len() {
            let card = self.code[self.pc];
            self.pc += 1;
            match card.rank {
                0u  => self.newline(card),
                K   => self.jump(card),
                Q   => continue,
                J   => self.dump(card, true),
                10u => self.dump(card, false),
                _   => self.assign(card)
            };
        }
        self
    }

    // Turn source file into array of cards.
    fn parse(filename: &str) -> Vec<Card> {
        let mut file = File::open(&Path::new(filename));
        let program = file.read_to_string().unwrap();

        // Split program blob into lines getting rid of comments and whitespaces.
        let comments = Regex::new(r"#.*$").unwrap();
        let lines: Vec<String> = program.as_slice().lines().map( |line|
            comments.replace_all(line, "").as_slice().trim().to_string()
        ).collect();

        // Turn source file into array of cards. Each card a struct of rank and suit.
        let mut code: Vec<Card> = vec![];
        let card = Regex::new(r"(10|[2-9JQKA])([♦♥♠♣])").unwrap();
        for (i, line) in lines.iter().enumerate() {
            // Line number cards have zero rank.
            code.push(Card { rank: 0, suit: i + 1 });

            // Parse cards using regural expression.
            for caps in card.captures_iter(line.as_slice()) {
                let rank = caps.at(1).char_at(0);
                let suit = caps.at(2).char_at(0);
                let card = match rank {
                   '1'       => Card { rank: 10, suit: suit as uint },
                   '2'...'9' => Card { rank: rank as uint - 48, suit: suit as uint },
                   _         => Card { rank: rank as uint, suit: suit as uint }
                };
                code.push(card);
            }
        }
        code
    }

    // Extra pass to set up jump labels.
    fn resolve(code: &Vec<Card>) -> HashMap<uint, uint> {
        let mut pc = 0u;
        let mut labels: HashMap<uint, uint> = HashMap::new();

        while pc < code.len() - 1 {
            let card = code[pc];
            pc += 1;
            if card.rank == Q {
                let mut queen = card.suit as uint;
                while pc < code.len() && code[pc].rank == Q && code[pc].suit == card.suit {
                    queen += card.suit as uint;
                    pc += 1;
                }
                labels.insert(queen, pc);
            }
        }
        labels
    }

    fn newline(&mut self, card: Card) -> &Ante {
        self.line = card.suit as uint;
        self
    }

    fn jump(&mut self, card: Card) -> &Ante {
        let mut suit = card.suit as uint;
        while self.pc < self.code.len() && self.code[self.pc].rank == K && self.code[self.pc].suit == card.suit {
            suit += card.suit as uint;
            self.pc += 1;
        }

        if !self.vars[card.suit].is_zero() {
            let label: uint = suit as uint;
            if self.labels.contains_key(&label) {
                self.pc = self.labels[label];
            } else {
                self.exception("can't find the label...");
            }
        }
        self
    }

    fn assign(&mut self, card: Card) -> &Ante {
        let operands = self.remaining(card);
        self.expression(operands)
    }

    fn dump(&mut self, card: Card, as_character: bool) -> &Ante {
        let value = self.vars[card.suit].clone();
        if as_character {
            if value < Ante::big(0) || value > Ante::big(255) {
                self.exception(format!("character code {} is not in 0..255 range", value).as_slice());
            } else {
                // Collect the bytes till we have full UTF-8 character.
                // Once the character is built dump it and reset the buffer.
                self.buffer.push(value.to_u8().unwrap());
                if is_utf8(self.buffer.as_slice()) {
                    print!("{}", from_utf8(self.buffer.as_slice()).unwrap());
                    self.buffer = vec![];
                }
            }
        } else {
            print!("{}", value);
        }
        self
    }

    fn remaining(&mut self, card: Card) -> Vec<Card> {
        let mut operands: Vec<Card> = vec![card];

        while self.pc < self.code.len() {
            let card = self.code[self.pc];
            if card.rank == 0 || card.rank == K || card.rank == Q || card.rank == J {
                break;
            }
            operands.push(card);
            self.pc += 1;
        }
        operands
    }

    fn expression(&mut self, operands: Vec<Card>) -> &Ante {
        let mut initial: BigInt = Ante::big(operands[0].rank as int);
        let target = operands[0].suit;

        if initial.to_uint().unwrap() == A {
            initial = self.vars[target].clone();
        }

        for i in range(1, operands.len()) {
            let mut rank = Ante::big(operands[i].rank as int);
            let suit = operands[i].suit;

            if rank.to_uint().unwrap() == A {
                rank = self.vars[suit].clone();
            }

            match suit {
                D => { initial = initial.add(&rank) },
                H => { initial = initial.mul(&rank) },
                S => { initial = initial.sub(&rank) },
                C => if !rank.is_zero() {
                        initial = initial.div(&rank);
                     } else {
                        self.exception("division by zero");
                     },
                _ => continue
            }
        }

        *self.vars.get_mut(&target) = initial;
        self
    }

    fn big(n: int) -> BigInt {
        from_int::<BigInt>(n).unwrap()
    }

    // NOTE: fail! got renamed to panic!
    fn exception(&self, message: &str) {
        fail!("Ante exception: {} on line {} (pc:{})\n", message, self.line, self.pc)
    }
}

fn main() {
    println!("usage: ante filename.ante");
    Ante::new("quine.ante".as_slice()).run();
}
<commit_msg>Load source file, handle file loading errors<commit_after>// Copyright (c) 2013-2014 Michael Dvorkin
// Ante is an esoteric programming language where all you've got is a deck of cards.
//
// This is Ante implementation in Rust.

extern crate regex;
extern crate num;

use std::io::File;
use std::collections::HashMap;
use std::num::Zero;
use std::num::from_int;
use std::str::{is_utf8, from_utf8};
use std::char::{from_u32};
use std::os;
use regex::Regex;
use num::bigint::BigInt;

const A: uint = 'A' as uint;
const J: uint = 'J' as uint;
const Q: uint = 'Q' as uint;
const K: uint = 'K' as uint;
const D: uint = '♦' as uint;
const H: uint = '♥' as uint;
const S: uint = '♠' as uint;
const C: uint = '♣' as uint;

struct Card {
    rank: uint,
    suit: uint
}

struct Ante {
    pc:     uint,                   // Program counter (index within ante.code)
    line:   uint,                   // Current line number.
    code:   Vec<Card>,              // Array of cards.
    vars:   HashMap<uint, BigInt>,  // Four registers hashed by suit.
    labels: HashMap<uint, uint>,    // Labels for ante.pc to jump to.
    buffer: Vec<u8>                 // Buffer to collect UTF-8 character bytes.
}

impl Ante {
    fn new(filename: &str) -> Ante {
        let code = Ante::parse(filename);
        let labels = Ante::resolve(&code);
        let mut vars: HashMap<uint, BigInt> = HashMap::new();
        vars.insert(D, Ante::big(0));
        vars.insert(H, Ante::big(0));
        vars.insert(S, Ante::big(0));
        vars.insert(C, Ante::big(0));

        Ante {
            pc:     0,
            line:   0,
            code:   code,
            vars:   vars,
            labels: labels,
            buffer: vec![]
        }
    }

    fn run(&mut self) -> &Ante {
        while self.pc < self.code.len() {
            let card = self.code[self.pc];
            self.pc += 1;
            match card.rank {
                0u  => self.newline(card),
                K   => self.jump(card),
                Q   => continue,
                J   => self.dump(card, true),
                10u => self.dump(card, false),
                _   => self.assign(card)
            };
        }
        self
    }

    // Turn source file into array of cards.
    fn parse(filename: &str) -> Vec<Card> {
        let mut file = match File::open(&Path::new(filename)) {
            // The `desc` field of `IoError` is a string that describes the error.
            Err(reason) => fail!("couldn't open {} ({})", filename, reason.desc),
            Ok(file) => file,
        };
        let program = file.read_to_string().unwrap();

        // Split program blob into lines getting rid of comments and whitespaces.
        let comments = Regex::new(r"#.*$").unwrap();
        let lines: Vec<String> = program.as_slice().lines().map( |line|
            comments.replace_all(line, "").as_slice().trim().to_string()
        ).collect();

        // Turn source file into array of cards. Each card a struct of rank and suit.
        let mut code: Vec<Card> = vec![];
        let card = Regex::new(r"(10|[2-9JQKA])([♦♥♠♣])").unwrap();
        for (i, line) in lines.iter().enumerate() {
            // Line number cards have zero rank.
            code.push(Card { rank: 0, suit: i + 1 });

            // Parse lines to turn them into array of cards.
            for caps in card.captures_iter(line.as_slice()) {
                let rank = caps.at(1).char_at(0);
                let suit = caps.at(2).char_at(0);
                let card = match rank {
                   '1'       => Card { rank: 10, suit: suit as uint },
                   '2'...'9' => Card { rank: rank as uint - 48, suit: suit as uint },
                   _         => Card { rank: rank as uint, suit: suit as uint }
                };
                code.push(card);
            }
        }
        code
    }

    // Extra pass to set up jump labels.
    fn resolve(code: &Vec<Card>) -> HashMap<uint, uint> {
        let mut pc = 0u;
        let mut labels: HashMap<uint, uint> = HashMap::new();

        while pc < code.len() - 1 {
            let card = code[pc];
            pc += 1;
            if card.rank == Q {
                let mut queen = card.suit as uint;
                while pc < code.len() && code[pc].rank == Q && code[pc].suit == card.suit {
                    queen += card.suit as uint;
                    pc += 1;
                }
                labels.insert(queen, pc);
            }
        }
        labels
    }

    fn newline(&mut self, card: Card) -> &Ante {
        self.line = card.suit as uint;
        self
    }

    fn jump(&mut self, card: Card) -> &Ante {
        let mut suit = card.suit as uint;
        while self.pc < self.code.len() && self.code[self.pc].rank == K && self.code[self.pc].suit == card.suit {
            suit += card.suit as uint;
            self.pc += 1;
        }

        if !self.vars[card.suit].is_zero() {
            let label = suit as uint;
            if self.labels.contains_key(&label) {
                self.pc = self.labels[label];
            } else {
                self.exception("can't find the label...");
            }
        }
        self
    }

    fn assign(&mut self, card: Card) -> &Ante {
        let operands = self.remaining(card);
        self.expression(operands)
    }

    fn dump(&mut self, card: Card, as_character: bool) -> &Ante {
        let value = self.vars[card.suit].clone();
        if as_character {
            if value < Ante::big(0) || value > Ante::big(255) {
                self.exception(format!("character code {} is not in 0..255 range", value).as_slice());
            } else {
                // Collect the bytes till we have full UTF-8 character.
                // Once the character is built dump it and reset the buffer.
                self.buffer.push(value.to_u8().unwrap());
                if is_utf8(self.buffer.as_slice()) {
                    print!("{}", from_utf8(self.buffer.as_slice()).unwrap());
                    self.buffer = vec![];
                }
            }
        } else {
            print!("{}", value);
        }
        self
    }

    fn remaining(&mut self, card: Card) -> Vec<Card> {
        let mut operands: Vec<Card> = vec![card];

        while self.pc < self.code.len() {
            let card = self.code[self.pc];
            if card.rank == 0 || card.rank == K || card.rank == Q || card.rank == J {
                break;
            }
            operands.push(card);
            self.pc += 1;
        }
        operands
    }

    fn expression(&mut self, operands: Vec<Card>) -> &Ante {
        let mut initial: BigInt = Ante::big(operands[0].rank as int);
        let target = operands[0].suit;

        if initial.to_uint().unwrap() == A {
            initial = self.vars[target].clone();
        }

        for i in range(1, operands.len()) {
            let mut rank = Ante::big(operands[i].rank as int);
            let suit = operands[i].suit;

            if rank.to_uint().unwrap() == A {
                rank = self.vars[suit].clone();
            }

            match suit {
                D => { initial = initial.add(&rank) },
                H => { initial = initial.mul(&rank) },
                S => { initial = initial.sub(&rank) },
                C => if !rank.is_zero() {
                        initial = initial.div(&rank);
                     } else {
                        self.exception("division by zero");
                     },
                _ => continue
            }
        }

        *self.vars.get_mut(&target) = initial;
        self
    }

    fn big(n: int) -> BigInt {
        from_int::<BigInt>(n).unwrap()
    }

    // NOTE: fail! got renamed to panic!
    fn exception(&self, message: &str) {
        fail!("Ante exception: {} on line {} (pc:{})\n", message, self.line, self.pc)
    }
}

fn main() {
    if os::args().len() == 2 {
        Ante::new(os::args()[1].as_slice()).run();
    } else {
        println!("usage: ante filename.ante");
    }
}
<|endoftext|>


<commit_before>use platform::vga::{LightRed, Black, Yellow};
use kernel::stdio::StdioWriter;

#[no_mangle]
pub fn entry()
{
	::platform::cpu::setup();
	main();
	loop { ::platform::cpu::halt(); }
}

fn main()
{
	let mut printer = StdioWriter::new();
	printer.bg = LightRed;
	printer.fg = Yellow;
	printer.clear_screen();
	printer.fg = Black;
	printer.move(3, 3);
	printer.print_screen("Hello, World!");
}
<commit_msg>Hello world is now more readable<commit_after>use platform::vga::{Red, White, Yellow};
use kernel::stdio::StdioWriter;

#[no_mangle]
pub fn entry()
{
	::platform::cpu::setup();
	main();
	loop { ::platform::cpu::halt(); }
}

fn main()
{
	let mut printer = StdioWriter::new();
	printer.bg = Red;
	printer.fg = Yellow;
	printer.clear_screen();
	printer.fg = White;
	printer.move(3, 3);
	printer.print_screen("Hello, World!");
}
<|endoftext|>


<commit_before>use engine::World;

pub struct Game {
    pub world: World,
}

impl Game {

    pub fn new() -> Game {
        Game {
            world: World::new(),
        }
    }

    pub fn step(&mut self) {

    }

}
<commit_msg>Add a log (MessageList) to the Game<commit_after>use engine::{MessageList, World};

pub struct Game {
    pub world: World,
    pub log: MessageList,
}

impl Game {

    pub fn new() -> Game {
        Game {
            world: World::new(),
            log: MessageList::new(),
        }
    }

    pub fn step(&mut self) {

    }

}
<|endoftext|>


<commit_before>#![recursion_limit = "256"]
#![feature(drain_filter)]
#![feature(test)]

pub mod client;
mod state;
mod stream;
pub mod utils;
pub mod wire;

pub use client::Client;
use state::State;
use stream::{Stream, StreamError};

use futures::{future::FutureExt, select, stream::StreamExt};
use std::{net::ToSocketAddrs, time::Duration};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::runtime::current_thread::Runtime;
use tokio::sync::mpsc;

//
// Public API
//

/// `Client` tries to reconnect on error after this many seconds.
pub const RECONNECT_SECS: u64 = 30;

#[derive(Debug, Clone)]
pub struct ServerInfo {
    /// Server address
    pub addr: String,

    /// Server port
    pub port: u16,

    /// Use TLS?
    pub tls: bool,

    /// Server password.
    pub pass: Option<String>,

    pub hostname: String,

    pub realname: String,

    /// Nicks to select when logging in.
    pub nicks: Vec<String>,

    /// Channels to automatically join
    pub auto_join: Vec<String>,

    /// Nickserv password. Sent to NickServ on connecting to the server and nick change, before
    /// join commands.
    pub nickserv_ident: Option<String>,
}

/// IRC client events. Returned by `Client` to the users via a channel.
#[derive(Debug)]
pub enum Event {
    /// Client trying to connect
    Connecting,
    /// TCP connection established *and* the introduction sequence with the IRC server started.
    Connected,
    /// Disconnected from the server. Usually sent right after an `Event::IoErr`.
    Disconnected,
    /// An IO error happened.
    IoErr(std::io::Error),
    /// A TLS error happened
    TlsErr(native_tls::Error),
    /// Client couldn't resolve host address.
    CantResolveAddr,
    /// Nick changed.
    NickChange(String),
    /// A message from the server
    Msg(wire::Msg),
}

impl From<StreamError> for Event {
    fn from(err: StreamError) -> Event {
        match err {
            StreamError::TlsError(tls_err) => Event::TlsErr(tls_err),
            StreamError::IoError(io_err) => Event::IoErr(io_err),
        }
    }
}

/// IRC client.
#[derive(Clone)]
pub struct IrcClient {
    /// Channel to the send commands to the main loop. Usually just for sending messages to the
    /// server.
    msg_chan: mpsc::Sender<Cmd>,

    // TODO: This is mostly here to make switching from the old `conn.rs` easier; it may be
    // possible to remove this and maybe have a unique usize in each client as id.
    serv_name: String,

    /// Reference to the state, to be able to provide methods like `get_nick` and
    /// `is_nick_accepted`.
    state: State,
    // We can't have a channel to the sender task directly here, because when the sender task
    // returns we lose the receiving end of the channel and there's no way to avoid this except
    // with annoying hacks like wrapping it with an `Arc<Mutex<..>>` or something.
}

impl IrcClient {
    /// Create a new client. Spawns two `tokio` tasks on the given `runtime`.
    pub fn new(
        server_info: ServerInfo,
        runtime: &mut Runtime,
    ) -> (IrcClient, mpsc::Receiver<Event>) {
        connect(server_info, runtime)
    }
}

impl Client for IrcClient {
    fn get_serv_name(&self) -> &str {
        &self.serv_name
    }

    // FIXME: This allocates a String
    fn get_nick(&self) -> String {
        self.state.get_nick()
    }

    // TODO: Do we really need this?
    fn is_nick_accepted(&self) -> bool {
        self.state.is_nick_accepted()
    }

    fn raw_msg(&mut self, msg: String) {
        self.msg_chan
            .try_send(Cmd::Msg(format!("{}\r\n", msg)))
            .unwrap()
    }

    fn split_privmsg<'a>(&self, extra_len: usize, msg: &'a str) -> utils::SplitIterator<'a> {
        // Max msg len calculation adapted from hexchat
        // (src/common/outbound.c:split_up_text)
        let mut max = 512; // RFC 2812
        max -= 3; // :, !, @
        max -= 13; // " PRIVMSG ", " ", :, \r, \n
        max -= self.get_nick().len();
        max -= extra_len;
        match self.state.get_usermask() {
            None => {
                max -= 9; // max username
                max -= 64; // max possible hostname (63) + '@'
                           // NOTE(osa): I think hexchat has an error here, it
                           // uses 65
            }
            Some(ref usermask) => {
                max -= usermask.len();
            }
        }

        assert!(max > 0);

        utils::split_iterator(msg, max)
    }

    fn privmsg(&mut self, target: &str, msg: &str, ctcp_action: bool) {
        let wire_fn = if ctcp_action {
            wire::privmsg
        } else {
            wire::action
        };
        self.msg_chan
            .try_send(Cmd::Msg(wire_fn(target, msg)))
            .unwrap()
    }

    fn join(&mut self, chans: &[&str]) {
        self.msg_chan
            .try_send(Cmd::Msg(wire::join(&chans)))
            .unwrap()
    }

    fn away(&mut self, msg: Option<&str>) {
        self.state.set_away(msg);
        self.msg_chan.try_send(Cmd::Msg(wire::away(msg))).unwrap()
    }

    fn nick(&mut self, new_nick: &str) {
        self.msg_chan
            .try_send(Cmd::Msg(wire::nick(new_nick)))
            .unwrap()
    }

    fn quit(&mut self, reason: Option<String>) {
        self.msg_chan.try_send(Cmd::Quit(reason)).unwrap();
    }
}

//
// End of public API
//

#[derive(Debug)]
enum Cmd {
    /// Send this IRC message to the server. Note that this needs to be a valid IRC message
    /// (including the trailing "\r\n").
    Msg(String),
    /// Close the connection. This sends a QUIT message to the server (with optional "reason") and
    /// then all tasks return.
    Quit(Option<String>),
}

fn connect(server_info: ServerInfo, runtime: &mut Runtime) -> (IrcClient, mpsc::Receiver<Event>) {
    let serv_name = server_info.addr.clone();

    //
    // Create communication channels
    //

    // Channel for returning IRC events to user.
    let (mut snd_ev, rcv_ev) = mpsc::channel::<Event>(100);
    // Channel for commands from user.
    let (snd_cmd, rcv_cmd) = mpsc::channel::<Cmd>(100);
    let mut rcv_cmd_fused = rcv_cmd.fuse();

    //
    // Create the main loop task
    //

    let irc_state = State::new(server_info.clone());
    let irc_state_clone = irc_state.clone();

    runtime.spawn(async move {
        // Main loop just tries to (re)connect
        'connect: loop {
            // Channel for the sender task. Messages are complete IRC messages (including the
            // trailing "\r\n") and the task directly sends them to the server.
            let (mut snd_msg, mut rcv_msg) = mpsc::channel::<String>(100);

            snd_ev.try_send(Event::Connecting).unwrap();

            //
            // Resolve IP address
            //

            let serv_name = server_info.addr.clone();
            let port = server_info.port;

            eprintln!("Resolving address");

            let serv_name_clone = serv_name.clone();
            let mut addr_iter = match tokio_executor::blocking::run(move || {
                (serv_name_clone.as_str(), port).to_socket_addrs()
            })
            .await
            {
                Err(io_err) => {
                    snd_ev.try_send(Event::IoErr(io_err)).unwrap();
                    tokio::timer::delay_for(Duration::from_secs(RECONNECT_SECS))
                        .await;
                    continue;
                }
                Ok(addr_iter) => addr_iter,
            };

            eprintln!("Address resolved");

            let addr = match addr_iter.next() {
                None => {
                    snd_ev.try_send(Event::CantResolveAddr).unwrap();
                    break;
                }
                Some(addr) => addr,
            };

            //
            // Establish TCP connection to the server
            //

            eprintln!("Establishing connection ...");

            let mb_stream = if server_info.tls {
                Stream::new_tls(addr, &serv_name).await
            } else {
                Stream::new_tcp(addr).await
            };

            let stream = match mb_stream {
                Ok(stream) => stream,
                Err(err) => {
                    snd_ev.try_send(Event::from(err)).unwrap();
                    snd_ev.try_send(Event::Disconnected).unwrap();
                    // Wait 30 seconds before looping
                    tokio::timer::delay_for(Duration::from_secs(RECONNECT_SECS)).await;
                    continue;
                }
            };

            let (mut read_half, mut write_half) = tokio::io::split(stream);

            eprintln!("Done");

            // let (mut read_half, mut write_half) = tokio::io::split(stream);

            //
            // Do the business
            //

            // Reset the connection state
            irc_state.reset();
            // Introduce self
            snd_msg.try_send(wire::nick(&irc_state.get_nick())).unwrap();
            snd_msg
                .try_send(wire::user(&server_info.hostname, &server_info.realname))
                .unwrap();

            // Spawn a task for outgoing messages.
            let mut snd_ev_clone = snd_ev.clone();
            tokio::spawn(async move {
                while let Some(msg) = rcv_msg.next().await {
                    if let Err(io_err) = write_half.write_all(msg.as_str().as_bytes()).await {
                        eprintln!("IO error when writing: {:?}", io_err);
                        snd_ev_clone.try_send(Event::IoErr(io_err)).unwrap();
                        return;
                    }
                }
            });

            let mut parse_buf: Vec<u8> = Vec::with_capacity(1024);

            // TODO: Introduce self here

            loop {
                let mut read_buf: [u8; 1024] = [0; 1024];

                select! {
                    cmd = rcv_cmd_fused.next() => {
                        match cmd {
                            None => {
                                eprintln!("main loop: command channel terminated from the other end");
                                // That's OK, rcv_cmd_fused will never be ready again
                            }
                            Some(Cmd::Msg(irc_msg)) => {
                                snd_msg.try_send(irc_msg).unwrap();
                            }
                            Some(Cmd::Quit(reason)) => {
                                snd_msg.send(wire::quit(reason)).await.unwrap();
                                // This drops the sender end of the channel that the sender task
                                // uses, which in turn causes the sender task to return. Somewhat
                                // hacky?
                                return;
                            }
                        }
                    }
                    bytes = read_half.read(&mut read_buf).fuse() => {
                        match bytes {
                            Err(io_err) => {
                                eprintln!("main loop: error when reading from socket: {:?}", io_err);
                                snd_ev.try_send(Event::IoErr(io_err)).unwrap();
                                snd_ev.try_send(Event::Disconnected).unwrap();
                                continue 'connect;
                            }
                            Ok(bytes) => {
                                parse_buf.extend_from_slice(&read_buf[0..bytes]);
                                while let Some(msg) = wire::parse_irc_msg(&mut parse_buf) {
                                    eprintln!("parsed msg: {:?}", msg);
                                    irc_state.update(&msg, &mut snd_ev, &mut snd_msg);
                                    snd_ev.try_send(Event::Msg(msg)).unwrap();
                                }
                            }
                        }
                    }
                }
            }
        }
    });

    (
        IrcClient {
            msg_chan: snd_cmd,
            serv_name,
            state: irc_state_clone,
        },
        rcv_ev,
    )
}
<commit_msg>Fix privmsg sending -- don't send ACTION by default<commit_after>#![recursion_limit = "256"]
#![feature(drain_filter)]
#![feature(test)]

pub mod client;
mod state;
mod stream;
pub mod utils;
pub mod wire;

pub use client::Client;
use state::State;
use stream::{Stream, StreamError};

use futures::{future::FutureExt, select, stream::StreamExt};
use std::{net::ToSocketAddrs, time::Duration};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::runtime::current_thread::Runtime;
use tokio::sync::mpsc;

//
// Public API
//

/// `Client` tries to reconnect on error after this many seconds.
pub const RECONNECT_SECS: u64 = 30;

#[derive(Debug, Clone)]
pub struct ServerInfo {
    /// Server address
    pub addr: String,

    /// Server port
    pub port: u16,

    /// Use TLS?
    pub tls: bool,

    /// Server password.
    pub pass: Option<String>,

    pub hostname: String,

    pub realname: String,

    /// Nicks to select when logging in.
    pub nicks: Vec<String>,

    /// Channels to automatically join
    pub auto_join: Vec<String>,

    /// Nickserv password. Sent to NickServ on connecting to the server and nick change, before
    /// join commands.
    pub nickserv_ident: Option<String>,
}

/// IRC client events. Returned by `Client` to the users via a channel.
#[derive(Debug)]
pub enum Event {
    /// Client trying to connect
    Connecting,
    /// TCP connection established *and* the introduction sequence with the IRC server started.
    Connected,
    /// Disconnected from the server. Usually sent right after an `Event::IoErr`.
    Disconnected,
    /// An IO error happened.
    IoErr(std::io::Error),
    /// A TLS error happened
    TlsErr(native_tls::Error),
    /// Client couldn't resolve host address.
    CantResolveAddr,
    /// Nick changed.
    NickChange(String),
    /// A message from the server
    Msg(wire::Msg),
}

impl From<StreamError> for Event {
    fn from(err: StreamError) -> Event {
        match err {
            StreamError::TlsError(tls_err) => Event::TlsErr(tls_err),
            StreamError::IoError(io_err) => Event::IoErr(io_err),
        }
    }
}

/// IRC client.
#[derive(Clone)]
pub struct IrcClient {
    /// Channel to the send commands to the main loop. Usually just for sending messages to the
    /// server.
    msg_chan: mpsc::Sender<Cmd>,

    // TODO: This is mostly here to make switching from the old `conn.rs` easier; it may be
    // possible to remove this and maybe have a unique usize in each client as id.
    serv_name: String,

    /// Reference to the state, to be able to provide methods like `get_nick` and
    /// `is_nick_accepted`.
    state: State,
    // We can't have a channel to the sender task directly here, because when the sender task
    // returns we lose the receiving end of the channel and there's no way to avoid this except
    // with annoying hacks like wrapping it with an `Arc<Mutex<..>>` or something.
}

impl IrcClient {
    /// Create a new client. Spawns two `tokio` tasks on the given `runtime`.
    pub fn new(
        server_info: ServerInfo,
        runtime: &mut Runtime,
    ) -> (IrcClient, mpsc::Receiver<Event>) {
        connect(server_info, runtime)
    }
}

impl Client for IrcClient {
    fn get_serv_name(&self) -> &str {
        &self.serv_name
    }

    // FIXME: This allocates a String
    fn get_nick(&self) -> String {
        self.state.get_nick()
    }

    // TODO: Do we really need this?
    fn is_nick_accepted(&self) -> bool {
        self.state.is_nick_accepted()
    }

    fn raw_msg(&mut self, msg: String) {
        self.msg_chan
            .try_send(Cmd::Msg(format!("{}\r\n", msg)))
            .unwrap()
    }

    fn split_privmsg<'a>(&self, extra_len: usize, msg: &'a str) -> utils::SplitIterator<'a> {
        // Max msg len calculation adapted from hexchat
        // (src/common/outbound.c:split_up_text)
        let mut max = 512; // RFC 2812
        max -= 3; // :, !, @
        max -= 13; // " PRIVMSG ", " ", :, \r, \n
        max -= self.get_nick().len();
        max -= extra_len;
        match self.state.get_usermask() {
            None => {
                max -= 9; // max username
                max -= 64; // max possible hostname (63) + '@'
                           // NOTE(osa): I think hexchat has an error here, it
                           // uses 65
            }
            Some(ref usermask) => {
                max -= usermask.len();
            }
        }

        assert!(max > 0);

        utils::split_iterator(msg, max)
    }

    fn privmsg(&mut self, target: &str, msg: &str, ctcp_action: bool) {
        let wire_fn = if ctcp_action {
            wire::action
        } else {
            wire::privmsg
        };
        self.msg_chan
            .try_send(Cmd::Msg(wire_fn(target, msg)))
            .unwrap()
    }

    fn join(&mut self, chans: &[&str]) {
        self.msg_chan
            .try_send(Cmd::Msg(wire::join(&chans)))
            .unwrap()
    }

    fn away(&mut self, msg: Option<&str>) {
        self.state.set_away(msg);
        self.msg_chan.try_send(Cmd::Msg(wire::away(msg))).unwrap()
    }

    fn nick(&mut self, new_nick: &str) {
        self.msg_chan
            .try_send(Cmd::Msg(wire::nick(new_nick)))
            .unwrap()
    }

    fn quit(&mut self, reason: Option<String>) {
        self.msg_chan.try_send(Cmd::Quit(reason)).unwrap();
    }
}

//
// End of public API
//

#[derive(Debug)]
enum Cmd {
    /// Send this IRC message to the server. Note that this needs to be a valid IRC message
    /// (including the trailing "\r\n").
    Msg(String),
    /// Close the connection. This sends a QUIT message to the server (with optional "reason") and
    /// then all tasks return.
    Quit(Option<String>),
}

fn connect(server_info: ServerInfo, runtime: &mut Runtime) -> (IrcClient, mpsc::Receiver<Event>) {
    let serv_name = server_info.addr.clone();

    //
    // Create communication channels
    //

    // Channel for returning IRC events to user.
    let (mut snd_ev, rcv_ev) = mpsc::channel::<Event>(100);
    // Channel for commands from user.
    let (snd_cmd, rcv_cmd) = mpsc::channel::<Cmd>(100);
    let mut rcv_cmd_fused = rcv_cmd.fuse();

    //
    // Create the main loop task
    //

    let irc_state = State::new(server_info.clone());
    let irc_state_clone = irc_state.clone();

    runtime.spawn(async move {
        // Main loop just tries to (re)connect
        'connect: loop {
            // Channel for the sender task. Messages are complete IRC messages (including the
            // trailing "\r\n") and the task directly sends them to the server.
            let (mut snd_msg, mut rcv_msg) = mpsc::channel::<String>(100);

            snd_ev.try_send(Event::Connecting).unwrap();

            //
            // Resolve IP address
            //

            let serv_name = server_info.addr.clone();
            let port = server_info.port;

            eprintln!("Resolving address");

            let serv_name_clone = serv_name.clone();
            let mut addr_iter = match tokio_executor::blocking::run(move || {
                (serv_name_clone.as_str(), port).to_socket_addrs()
            })
            .await
            {
                Err(io_err) => {
                    snd_ev.try_send(Event::IoErr(io_err)).unwrap();
                    tokio::timer::delay_for(Duration::from_secs(RECONNECT_SECS))
                        .await;
                    continue;
                }
                Ok(addr_iter) => addr_iter,
            };

            eprintln!("Address resolved");

            let addr = match addr_iter.next() {
                None => {
                    snd_ev.try_send(Event::CantResolveAddr).unwrap();
                    break;
                }
                Some(addr) => addr,
            };

            //
            // Establish TCP connection to the server
            //

            eprintln!("Establishing connection ...");

            let mb_stream = if server_info.tls {
                Stream::new_tls(addr, &serv_name).await
            } else {
                Stream::new_tcp(addr).await
            };

            let stream = match mb_stream {
                Ok(stream) => stream,
                Err(err) => {
                    snd_ev.try_send(Event::from(err)).unwrap();
                    snd_ev.try_send(Event::Disconnected).unwrap();
                    // Wait 30 seconds before looping
                    tokio::timer::delay_for(Duration::from_secs(RECONNECT_SECS)).await;
                    continue;
                }
            };

            let (mut read_half, mut write_half) = tokio::io::split(stream);

            eprintln!("Done");

            // let (mut read_half, mut write_half) = tokio::io::split(stream);

            //
            // Do the business
            //

            // Reset the connection state
            irc_state.reset();
            // Introduce self
            snd_msg.try_send(wire::nick(&irc_state.get_nick())).unwrap();
            snd_msg
                .try_send(wire::user(&server_info.hostname, &server_info.realname))
                .unwrap();

            // Spawn a task for outgoing messages.
            let mut snd_ev_clone = snd_ev.clone();
            tokio::spawn(async move {
                while let Some(msg) = rcv_msg.next().await {
                    if let Err(io_err) = write_half.write_all(msg.as_str().as_bytes()).await {
                        eprintln!("IO error when writing: {:?}", io_err);
                        snd_ev_clone.try_send(Event::IoErr(io_err)).unwrap();
                        return;
                    }
                }
            });

            let mut parse_buf: Vec<u8> = Vec::with_capacity(1024);

            // TODO: Introduce self here

            loop {
                let mut read_buf: [u8; 1024] = [0; 1024];

                select! {
                    cmd = rcv_cmd_fused.next() => {
                        match cmd {
                            None => {
                                eprintln!("main loop: command channel terminated from the other end");
                                // That's OK, rcv_cmd_fused will never be ready again
                            }
                            Some(Cmd::Msg(irc_msg)) => {
                                snd_msg.try_send(irc_msg).unwrap();
                            }
                            Some(Cmd::Quit(reason)) => {
                                snd_msg.send(wire::quit(reason)).await.unwrap();
                                // This drops the sender end of the channel that the sender task
                                // uses, which in turn causes the sender task to return. Somewhat
                                // hacky?
                                return;
                            }
                        }
                    }
                    bytes = read_half.read(&mut read_buf).fuse() => {
                        match bytes {
                            Err(io_err) => {
                                eprintln!("main loop: error when reading from socket: {:?}", io_err);
                                snd_ev.try_send(Event::IoErr(io_err)).unwrap();
                                snd_ev.try_send(Event::Disconnected).unwrap();
                                continue 'connect;
                            }
                            Ok(bytes) => {
                                parse_buf.extend_from_slice(&read_buf[0..bytes]);
                                while let Some(msg) = wire::parse_irc_msg(&mut parse_buf) {
                                    eprintln!("parsed msg: {:?}", msg);
                                    irc_state.update(&msg, &mut snd_ev, &mut snd_msg);
                                    snd_ev.try_send(Event::Msg(msg)).unwrap();
                                }
                            }
                        }
                    }
                }
            }
        }
    });

    (
        IrcClient {
            msg_chan: snd_cmd,
            serv_name,
            state: irc_state_clone,
        },
        rcv_ev,
    )
}
<|endoftext|>