<commit_before>use luminance::{FragmentShader, GeometryShader, StageError, ShaderTypeable,
                TessellationControlShader, TessellationEvaluationShader, VertexShader};
use luminance_gl::gl33::{ProgramProxy, Stage};
use std::fs;
use std::io::Read;
use std::path::Path;

pub use luminance::{ProgramError, Uniformable, UniformUpdate};
pub use luminance::shader::program::UniformWarning;
pub use luminance_gl::gl33::{self, Uniform};
pub use luminance_gl::gl33::token::*;

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

fn compile_stages(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str) -> Result<(Option<(Stage<TessellationControlShader>, Stage<TessellationEvaluationShader>)>, Stage<VertexShader>, Option<Stage<GeometryShader>>, Stage<FragmentShader>), StageError> {
  let tess = match tess_src {
    None => None,
    Some((tcs_src, tes_src)) => {
      let tcs = try!(Stage::new(tcs_src));
      let tes = try!(Stage::new(tes_src));
      Some((tcs, tes))
    }
  };
  let vs = try!(Stage::new(vs_src));
  let gs = match gs_src {
    None => None,
    Some(gs_src) => {
      let gs = try!(Stage::new(gs_src));
      Some(gs)
    }
  };
  let fs = try!(Stage::new(fs_src));

  Ok((tess, vs, gs, fs))
}

pub fn new_program<GetUni, T>(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str, get_uni: &GetUni) -> Result<gl33::Program<T>, ProgramError> where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> {
  let stages = compile_stages(tess_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          gl33::Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
        } else {
          gl33::Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
        }
      } else if let Some(gs) = gs {
        gl33::Program::new(None, &vs, Some(&gs), &fs, get_uni)
      } else {
        gl33::Program::new(None, &vs, None, &fs, get_uni)
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

pub fn new_program_from_disk<P, GetUni, T>(tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: &GetUni) -> Result<gl33::Program<T>, ProgramError>
      where P: AsRef<Path>,
            GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> {
  // load vertex and fragment shaders first
  let vs = try!(read_stage(vs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
  let fs = try!(read_stage(fs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));

  match (tess_path, gs_path) {
    (None, None) => { // no tessellation nor geometry
      gl33::Program::new(None, &vs, None, &fs, get_uni)
    },
    (Some((tcs_path, tes_path)), None) => { // tessellation without geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
    },
    (None, Some(gs_path)) => { // geometry without tessellation
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(None, &vs, Some(&gs), &fs, get_uni)
    },
    (Some((tcs_path, tes_path)), Some(gs_path)) => { // tessellation and geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
    }
  }
}

pub fn read_stage<P, T>(path: P) -> Result<Stage<T>, StageError> where P: AsRef<Path>, T: ShaderTypeable {
  let path = path.as_ref();

  info!("loading {:?} stage: \x1b[35m{:?}", T::shader_type(), path);

  let fh = fs::File::open(path);

  match fh {
    Err(e) => {
      Err(StageError::CompilationFailed(T::shader_type(), format!("{:?}", e)))
    },
    Ok(mut fh) => {
      let mut stage_src = String::new();
      let _ = fh.read_to_string(&mut stage_src);

      Stage::new(stage_src.chars().as_str())
    }
  }
}

#[cfg(feature = "hot-resource")]
mod hot {
  use luminance_gl::gl33::ProgramProxy;
  use std::ops::Deref;
  use std::path::{Path, PathBuf};
  use std::sync::mpsc;

  use resource::ResourceManager;

  use super::*;

  pub struct Program<T> {
    rx: mpsc::Receiver<()>,
    last_update_time: Option<f64>,
    program: gl33::Program<T>,
    get_uni: Box<Fn(ProgramProxy) -> Result<T, ProgramError>>,
    vs_path: PathBuf,
    fs_path: PathBuf,
    tess_path: Option<(PathBuf, PathBuf)>,
    gs_path: Option<PathBuf>
  }

  impl<T> Program<T> {
    pub fn load<GetUni, P>(manager: &mut ResourceManager, tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: GetUni) -> Result<Self, ProgramError>
        where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> + 'static,
              P: AsRef<Path> {
      let vs_path = vs_path.as_ref();
      let fs_path = fs_path.as_ref();
      let tess_path = tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_ref(), tes.as_ref()));
      let gs_path = gs_path.as_ref().map(|gs| gs.as_ref());

      let program = try!(new_program_from_disk(tess_path, vs_path, gs_path, fs_path, &get_uni));

      let (sx, rx) = mpsc::channel();
      manager.monitor(vs_path, sx.clone());
      manager.monitor(fs_path, sx.clone());

      if let Some((tcs, tes)) = tess_path {
        manager.monitor(tcs, sx.clone());
        manager.monitor(tes, sx.clone());
      }

      if let Some(gs) = gs_path {
        manager.monitor(gs, sx.clone());
      }

      Ok(Program {
        rx: rx,
        last_update_time: None,
        program: program,
        get_uni: Box::new(get_uni),
        vs_path: vs_path.to_path_buf(),
        fs_path: fs_path.to_path_buf(),
        tess_path: tess_path.map(|(tcs, tes)| (tcs.to_path_buf(), tes.to_path_buf())),
        gs_path: gs_path.map(|gs| gs.to_path_buf()),
      })
    }

    fn reload(&mut self) {
      let vs = self.vs_path.as_path();
      let fs = self.fs_path.as_path();
      let tess = self.tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_path(), tes.as_path()));
      let gs = self.gs_path.as_ref().map(|gs| gs.as_path());

      let program = new_program_from_disk(tess, vs, gs, fs, &self.get_uni.as_ref());

      match program {
        Ok(program) => {
          self.program = program;
          info!("reloaded shader program");
        },
        Err(e) => {
          err!("reloading program has failed: {:?}", e);
        }
      }
    }

    decl_sync_hot!();
  }

  impl<'a, T> Deref for Program<T> {
    type Target = gl33::Program<T>;

    fn deref(&self) -> &Self::Target {
      &self.program
    }
  }
}

#[cfg(not(feature = "hot-resource"))]
mod cold {
  use super::*;

  use luminance_gl::gl33::ProgramProxy;
  use resource::ResourceManager;
  use std::ops::Deref;
  use std::path::Path;

  pub struct Program<T>(gl33::Program<T>);

  impl<T> Program<T> {
    pub fn load<GetUni, P>(_: &mut ResourceManager, tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: GetUni) -> Result<Program<T>, ProgramError>
        where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> + 'static,
              P: AsRef<Path> {
      let vs_path = vs_path.as_ref();
      let fs_path = fs_path.as_ref();
      let tess_path = tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_ref(), tes.as_ref()));
      let gs_path = gs_path.as_ref().map(|gs| gs.as_ref());

      let program = try!(new_program_from_disk(tess_path, vs_path, gs_path, fs_path, &get_uni));

      Ok(Program(program))
    }

    pub fn sync(&mut self) {}
  }

  impl<'a, T> Deref for Program<T> {
    type Target = gl33::Program<T>;

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

#[cfg(feature = "hot-resource")]
pub use self::hot::*;
#[cfg(not(feature = "hot-resource"))]
pub use self::cold::*;

/// A helper function used to make uniforms optionable. If there’s a warning, it’s printed out.
pub fn opt_uni<T>(uni: (Uniform<T>, Option<UniformWarning>)) -> Uniform<T> where T: Uniformable<GL33> {
  if let Some(warning) = uni.1 {
    warn!("{:?}", warning);
  }

  uni.0
}
<commit_msg>Fixed ProgramError -> UniformWarning for uniform interface builder function.<commit_after>use luminance::{FragmentShader, GeometryShader, StageError, ShaderTypeable,
                TessellationControlShader, TessellationEvaluationShader, VertexShader};
use luminance_gl::gl33::{ProgramProxy, Stage};
use std::fs;
use std::io::Read;
use std::path::Path;

pub use luminance::{ProgramError, Uniformable, UniformUpdate};
pub use luminance::shader::program::UniformWarning;
pub use luminance_gl::gl33::{self, Uniform};
pub use luminance_gl::gl33::token::*;

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

fn compile_stages(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str) -> Result<(Option<(Stage<TessellationControlShader>, Stage<TessellationEvaluationShader>)>, Stage<VertexShader>, Option<Stage<GeometryShader>>, Stage<FragmentShader>), StageError> {
  let tess = match tess_src {
    None => None,
    Some((tcs_src, tes_src)) => {
      let tcs = try!(Stage::new(tcs_src));
      let tes = try!(Stage::new(tes_src));
      Some((tcs, tes))
    }
  };
  let vs = try!(Stage::new(vs_src));
  let gs = match gs_src {
    None => None,
    Some(gs_src) => {
      let gs = try!(Stage::new(gs_src));
      Some(gs)
    }
  };
  let fs = try!(Stage::new(fs_src));

  Ok((tess, vs, gs, fs))
}

pub fn new_program<GetUni, T>(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str, get_uni: &GetUni) -> Result<gl33::Program<T>, ProgramError> where GetUni: Fn(ProgramProxy) -> Result<T, UniformWarning> {
  let stages = compile_stages(tess_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          gl33::Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
        } else {
          gl33::Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
        }
      } else if let Some(gs) = gs {
        gl33::Program::new(None, &vs, Some(&gs), &fs, get_uni)
      } else {
        gl33::Program::new(None, &vs, None, &fs, get_uni)
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

pub fn new_program_from_disk<P, GetUni, T>(tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: &GetUni) -> Result<gl33::Program<T>, ProgramError>
      where P: AsRef<Path>,
            GetUni: Fn(ProgramProxy) -> Result<T, UniformWarning> {
  // load vertex and fragment shaders first
  let vs = try!(read_stage(vs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
  let fs = try!(read_stage(fs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));

  match (tess_path, gs_path) {
    (None, None) => { // no tessellation nor geometry
      gl33::Program::new(None, &vs, None, &fs, get_uni)
    },
    (Some((tcs_path, tes_path)), None) => { // tessellation without geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
    },
    (None, Some(gs_path)) => { // geometry without tessellation
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(None, &vs, Some(&gs), &fs, get_uni)
    },
    (Some((tcs_path, tes_path)), Some(gs_path)) => { // tessellation and geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      gl33::Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
    }
  }
}

pub fn read_stage<P, T>(path: P) -> Result<Stage<T>, StageError> where P: AsRef<Path>, T: ShaderTypeable {
  let path = path.as_ref();

  info!("loading {:?} stage: \x1b[35m{:?}", T::shader_type(), path);

  let fh = fs::File::open(path);

  match fh {
    Err(e) => {
      Err(StageError::CompilationFailed(T::shader_type(), format!("{:?}", e)))
    },
    Ok(mut fh) => {
      let mut stage_src = String::new();
      let _ = fh.read_to_string(&mut stage_src);

      Stage::new(stage_src.chars().as_str())
    }
  }
}

#[cfg(feature = "hot-resource")]
mod hot {
  use luminance_gl::gl33::ProgramProxy;
  use std::ops::Deref;
  use std::path::{Path, PathBuf};
  use std::sync::mpsc;

  use resource::ResourceManager;

  use super::*;

  pub struct Program<T> {
    rx: mpsc::Receiver<()>,
    last_update_time: Option<f64>,
    program: gl33::Program<T>,
    get_uni: Box<Fn(ProgramProxy) -> Result<T, UniformWarning>>,
    vs_path: PathBuf,
    fs_path: PathBuf,
    tess_path: Option<(PathBuf, PathBuf)>,
    gs_path: Option<PathBuf>
  }

  impl<T> Program<T> {
    pub fn load<GetUni, P>(manager: &mut ResourceManager, tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: GetUni) -> Result<Self, ProgramError>
        where GetUni: Fn(ProgramProxy) -> Result<T, UniformWarning> + 'static,
              P: AsRef<Path> {
      let vs_path = vs_path.as_ref();
      let fs_path = fs_path.as_ref();
      let tess_path = tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_ref(), tes.as_ref()));
      let gs_path = gs_path.as_ref().map(|gs| gs.as_ref());

      let program = try!(new_program_from_disk(tess_path, vs_path, gs_path, fs_path, &get_uni));

      let (sx, rx) = mpsc::channel();
      manager.monitor(vs_path, sx.clone());
      manager.monitor(fs_path, sx.clone());

      if let Some((tcs, tes)) = tess_path {
        manager.monitor(tcs, sx.clone());
        manager.monitor(tes, sx.clone());
      }

      if let Some(gs) = gs_path {
        manager.monitor(gs, sx.clone());
      }

      Ok(Program {
        rx: rx,
        last_update_time: None,
        program: program,
        get_uni: Box::new(get_uni),
        vs_path: vs_path.to_path_buf(),
        fs_path: fs_path.to_path_buf(),
        tess_path: tess_path.map(|(tcs, tes)| (tcs.to_path_buf(), tes.to_path_buf())),
        gs_path: gs_path.map(|gs| gs.to_path_buf()),
      })
    }

    fn reload(&mut self) {
      let vs = self.vs_path.as_path();
      let fs = self.fs_path.as_path();
      let tess = self.tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_path(), tes.as_path()));
      let gs = self.gs_path.as_ref().map(|gs| gs.as_path());

      let program = new_program_from_disk(tess, vs, gs, fs, &self.get_uni.as_ref());

      match program {
        Ok(program) => {
          self.program = program;
          info!("reloaded shader program");
        },
        Err(e) => {
          err!("reloading program has failed: {:?}", e);
        }
      }
    }

    decl_sync_hot!();
  }

  impl<'a, T> Deref for Program<T> {
    type Target = gl33::Program<T>;

    fn deref(&self) -> &Self::Target {
      &self.program
    }
  }
}

#[cfg(not(feature = "hot-resource"))]
mod cold {
  use super::*;

  use luminance_gl::gl33::ProgramProxy;
  use resource::ResourceManager;
  use std::ops::Deref;
  use std::path::Path;

  pub struct Program<T>(gl33::Program<T>);

  impl<T> Program<T> {
    pub fn load<GetUni, P>(_: &mut ResourceManager, tess_path: Option<(P, P)>, vs_path: P, gs_path: Option<P>, fs_path: P, get_uni: GetUni) -> Result<Program<T>, ProgramError>
        where GetUni: Fn(ProgramProxy) -> Result<T, UniformWarning> + 'static,
              P: AsRef<Path> {
      let vs_path = vs_path.as_ref();
      let fs_path = fs_path.as_ref();
      let tess_path = tess_path.as_ref().map(|&(ref tcs, ref tes)| (tcs.as_ref(), tes.as_ref()));
      let gs_path = gs_path.as_ref().map(|gs| gs.as_ref());

      let program = try!(new_program_from_disk(tess_path, vs_path, gs_path, fs_path, &get_uni));

      Ok(Program(program))
    }

    pub fn sync(&mut self) {}
  }

  impl<'a, T> Deref for Program<T> {
    type Target = gl33::Program<T>;

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

#[cfg(feature = "hot-resource")]
pub use self::hot::*;
#[cfg(not(feature = "hot-resource"))]
pub use self::cold::*;

/// A helper function used to make uniforms optionable. If there’s a warning, it’s printed out.
pub fn opt_uni<T>(uni: (Uniform<T>, Option<UniformWarning>)) -> Uniform<T> where T: Uniformable<GL33> {
  if let Some(warning) = uni.1 {
    warn!("{:?}", warning);
  }

  uni.0
}
<|endoftext|>


<commit_before>use luminance::{FragmentShader, GeometryShader, StageError, ShaderTypeable,
                TessellationControlShader, TessellationEvaluationShader, VertexShader};
use luminance_gl::gl33::{ProgramProxy, Stage};
use notify::{self, RecommendedWatcher, Watcher};
use std::collections::BTreeMap;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::thread;

pub use luminance::{ProgramError, UniformUpdate};
use luminance::shader::stage;
pub use luminance_gl::gl33::{Program, Uniform};

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

fn compile_stages(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str) -> Result<(Option<(Stage<TessellationControlShader>, Stage<TessellationEvaluationShader>)>, Stage<VertexShader>, Option<Stage<GeometryShader>>, Stage<FragmentShader>), StageError> {
  let tess = match tess_src {
    None => None,
    Some((tcs_src, tes_src)) => {
      let tcs = try!(Stage::new(tcs_src));
      let tes = try!(Stage::new(tes_src));
      Some((tcs, tes))
    }
  };
  let vs = try!(Stage::new(vs_src));
  let gs = match gs_src {
    None => None,
    Some(gs_src) => {
      let gs = try!(Stage::new(gs_src));
      Some(gs)
    }
  };
  let fs = try!(Stage::new(fs_src));

  Ok((tess, vs, gs, fs))
}

pub fn new_program<GetUni, T>(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str, get_uni: GetUni) -> Result<Program<T>, ProgramError> where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> {
  let stages = compile_stages(tess_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
        } else {
          Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
        }
      } else {
        if let Some(gs) = gs {
          Program::new(None, &vs, Some(&gs), &fs, get_uni)
        } else {
          Program::new(None, &vs, None, &fs, get_uni)
        }
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

pub fn new_program_from_disk<GetUni, T>(tess_path: Option<(&Path, &Path)>, vs_path: &Path, gs_path: Option<&Path>, fs_path: &Path, get_uni: GetUni) -> Result<Program<T>, ProgramError>
      where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> + Clone {
  // load vertex and fragment shaders first
  let vs = try!(read_stage(vs_path.as_ref()).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
  let fs = try!(read_stage(fs_path.as_ref()).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));

  match (&tess_path, &gs_path) {
    (&None, &None) => { // no tessellation nor geometry
      Program::new(None, &vs, None, &fs, get_uni.clone())
    },
    (&Some((ref tcs_path, ref tes_path)), &None) => { // tessellation without geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni.clone())
    },
    (&None, &Some(ref gs_path)) => { // geometry without tessellation
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(None, &vs, Some(&gs), &fs, get_uni.clone())
    },
    (&Some((ref tcs_path, ref tes_path)), &Some(ref gs_path)) => { // tessellation and geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni.clone())
    }
  }
}

pub fn read_stage<T>(path: &Path) -> Result<Stage<T>, StageError> where T: ShaderTypeable {
  info!("\tloading {:?} stage: \x1b[35m{:?}", T::shader_type(), path);

  let fh = fs::File::open(path);

  match fh {
    Err(e) => {
      Err(StageError::CompilationFailed(T::shader_type(), format!("{:?}", e)))
    },
    Ok(mut fh) => {
      let mut stage_src = String::new();
      let _ = fh.read_to_string(&mut stage_src);

      Stage::new(stage_src.chars().as_str())
    }
  }
}

/// Add surveillance of a given `Program` by providing the path to all its shaders. When a change
/// occurs, the `Program` gets notified of the change via its `Receiver` channel part.
pub fn monitor_shader(tess: Option<(&Path, &Path)>, vs: &Path, gs: Option<&Path>, fs: &Path, sx: mpsc::Sender<()>) {
  let (wsx, wrx) = mpsc::channel();
  let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

  // vertex shader
  watcher.watch(vs);

  // fragment shader
  watcher.watch(fs);

  // tessellation, if needed
  if let Some((tcs, tes)) = tess {
    // tessellation control shader
    watcher.watch(tcs);
    // tessellation evaluation shader
    watcher.watch(tes);
  }

  // geometry shader, if needed
  if let Some(gs) = gs {
    watcher.watch(gs);
  }

  // start a new monitoring thread
  let _ = thread::spawn(move || {
    loop {
      if let Ok(notify::Event { path: Some(path), op: Ok(notify::op::WRITE) }) = wrx.recv() {
        deb!("{:?}’s content has changed!", path);
        sx.send(());
      }
    }
  });
}

/// A `Program` wrapped by **ion**.
///
/// That wrapper is used to enable hot-reloading of shader programs.
pub struct WrappedProgram<'a, T> {
  rx: mpsc::Receiver<()>,
  program: Program<T>,
  get_uni: Box<Fn(ProgramProxy) -> Result<T, ProgramError> + 'a>,
  vs_path: PathBuf,
  fs_path: PathBuf,
  tess_path: Option<(PathBuf, PathBuf)>,
  gs_path: Option<PathBuf>
}

impl<'a, T> WrappedProgram<'a, T> {
  pub fn new<GetUni>(tess_path: Option<(&Path, &Path)>, vs_path: &Path, gs_path: Option<&Path>, fs_path: &Path, get_uni: GetUni) -> Result<Self, ProgramError>
      where GetUni: 'a + Fn(ProgramProxy) -> Result<T, ProgramError> + Clone {
    let program = try!(new_program_from_disk(tess_path, vs_path, gs_path, fs_path, get_uni.clone()));
    let (sx, rx) = mpsc::channel();

    monitor_shader(tess_path, vs_path, gs_path, fs_path, sx);

    let wrapped = WrappedProgram {
      rx: rx,
      program: program,
      get_uni: Box::new(get_uni),
      vs_path: vs_path.to_path_buf(),
      fs_path: fs_path.to_path_buf(),
      tess_path: tess_path.map(|(tcs, tes)| (tcs.to_path_buf(), tes.to_path_buf())),
      gs_path: gs_path.map(|gs| gs.to_path_buf())
    };

    Ok(wrapped)
  }

  //fn reload(&self mut) {
  //  let program =
  //}

  /// Sync the embedded `Program`.
  pub fn sync(&mut self) {
    if self.rx.try_recv().is_ok() {
    }
  }
}
<commit_msg>Ditto.<commit_after>use luminance::{FragmentShader, GeometryShader, StageError, ShaderTypeable,
                TessellationControlShader, TessellationEvaluationShader, VertexShader};
use luminance_gl::gl33::{ProgramProxy, Stage};
use notify::{self, RecommendedWatcher, Watcher};
use std::fs;
use std::io::Read;
use std::path::PathBuf;
use std::sync::mpsc;
use std::thread;

pub use luminance::{ProgramError, UniformUpdate};
use luminance::shader::stage;
pub use luminance_gl::gl33::{Program, Uniform};

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

fn compile_stages(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str) -> Result<(Option<(Stage<TessellationControlShader>, Stage<TessellationEvaluationShader>)>, Stage<VertexShader>, Option<Stage<GeometryShader>>, Stage<FragmentShader>), StageError> {
  let tess = match tess_src {
    None => None,
    Some((tcs_src, tes_src)) => {
      let tcs = try!(Stage::new(tcs_src));
      let tes = try!(Stage::new(tes_src));
      Some((tcs, tes))
    }
  };
  let vs = try!(Stage::new(vs_src));
  let gs = match gs_src {
    None => None,
    Some(gs_src) => {
      let gs = try!(Stage::new(gs_src));
      Some(gs)
    }
  };
  let fs = try!(Stage::new(fs_src));

  Ok((tess, vs, gs, fs))
}

pub fn new_program<GetUni, T>(tess_src: Option<(&str, &str)>, vs_src: &str, gs_src: Option<&str>, fs_src: &str, get_uni: GetUni) -> Result<Program<T>, ProgramError> where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> {
  let stages = compile_stages(tess_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni)
        } else {
          Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni)
        }
      } else {
        if let Some(gs) = gs {
          Program::new(None, &vs, Some(&gs), &fs, get_uni)
        } else {
          Program::new(None, &vs, None, &fs, get_uni)
        }
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

pub fn new_program_from_disk<GetUni, T>(tess_path: Option<(PathBuf, PathBuf)>, vs_path: PathBuf, gs_path: Option<PathBuf>, fs_path: PathBuf, get_uni: GetUni) -> Result<Program<T>, ProgramError>
      where GetUni: Fn(ProgramProxy) -> Result<T, ProgramError> + Clone {
  // load vertex and fragment shaders first
  let vs = try!(read_stage(vs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
  let fs = try!(read_stage(fs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));

  match (tess_path, gs_path) {
    (None, None) => { // no tessellation nor geometry
      Program::new(None, &vs, None, &fs, get_uni.clone())
    },
    (Some((tcs_path, tes_path)), None) => { // tessellation without geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(Some((&tcs, &tes)), &vs, None, &fs, get_uni.clone())
    },
    (None, Some(gs_path)) => { // geometry without tessellation
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(None, &vs, Some(&gs), &fs, get_uni.clone())
    },
    (Some((tcs_path, tes_path)), Some(gs_path)) => { // tessellation and geometry
      let tcs = try!(read_stage(tcs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let tes = try!(read_stage(tes_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      let gs = try!(read_stage(gs_path).map_err(|e| ProgramError::LinkFailed(format!("{:?}", e))));
      Program::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs, get_uni.clone())
    }
  }
}

pub fn read_stage<T>(path: PathBuf) -> Result<Stage<T>, StageError> where T: ShaderTypeable {
  info!("\tloading {:?} stage: \x1b[35m{:?}", T::shader_type(), path);

  let fh = fs::File::open(path);

  match fh {
    Err(e) => {
      Err(StageError::CompilationFailed(T::shader_type(), format!("{:?}", e)))
    },
    Ok(mut fh) => {
      let mut stage_src = String::new();
      let _ = fh.read_to_string(&mut stage_src);

      Stage::new(stage_src.chars().as_str())
    }
  }
}

/// Add surveillance of a given `Program` by providing the path to all its shaders. When a change
/// occurs, the `Program` gets notified of the change via its `Receiver` channel part.
pub fn monitor_shader(tess: Option<(PathBuf, PathBuf)>, vs: PathBuf, gs: Option<PathBuf>, fs: PathBuf, sx: mpsc::Sender<()>) {
  let (wsx, wrx) = mpsc::channel();
  let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

  // vertex shader
  watcher.watch(vs);

  // fragment shader
  watcher.watch(fs);

  // tessellation, if needed
  if let Some((tcs, tes)) = tess {
    // tessellation control shader
    watcher.watch(tcs);
    // tessellation evaluation shader
    watcher.watch(tes);
  }

  // geometry shader, if needed
  if let Some(gs) = gs {
    watcher.watch(gs);
  }

  // start a new monitoring thread
  let _ = thread::spawn(move || {
    loop {
      if let Ok(notify::Event { path: Some(path), op: Ok(notify::op::WRITE) }) = wrx.recv() {
        deb!("{:?}’s content has changed!", path);
        sx.send(());
      }
    }
  });
}

/// A `Program` wrapped by **ion**.
///
/// That wrapper is used to enable hot-reloading of shader programs.
pub struct WrappedProgram<'a, T> {
  rx: mpsc::Receiver<()>,
  program: Program<T>,
  get_uni: Box<Fn(ProgramProxy) -> Result<T, ProgramError> + 'a>,
  vs_path: PathBuf,
  fs_path: PathBuf,
  tess_path: Option<(PathBuf, PathBuf)>,
  gs_path: Option<PathBuf>
}

impl<'a, T> WrappedProgram<'a, T> {
  pub fn new<GetUni>(tess_path: Option<(PathBuf, PathBuf)>, vs_path: PathBuf, gs_path: Option<PathBuf>, fs_path: PathBuf, get_uni: GetUni) -> Result<Self, ProgramError>
      where GetUni: 'a + Fn(ProgramProxy) -> Result<T, ProgramError> + Clone {
    let program = try!(new_program_from_disk(tess_path.clone(), vs_path.clone(), gs_path.clone(), fs_path.clone(), get_uni.clone()));
    let (sx, rx) = mpsc::channel();

    monitor_shader(tess_path.clone(), vs_path.clone(), gs_path.clone(), fs_path.clone(), sx);

    let wrapped = WrappedProgram {
      rx: rx,
      program: program,
      get_uni: Box::new(get_uni),
      vs_path: vs_path,
      fs_path: fs_path,
      tess_path: tess_path,
      gs_path: gs_path,
    };

    Ok(wrapped)
  }

  //fn reload(&self mut) {
  //  let program = new_program_from_disk(self.tess_path,
  //}

  /// Sync the embedded `Program`.
  pub fn sync(&mut self) {
    if self.rx.try_recv().is_ok() {
    }
  }
}
<|endoftext|>


<commit_before><commit_msg>Rust coding style<commit_after><|endoftext|>


<commit_before>use super::*;
use super::super::packets::*;
use super::super::async;
use super::super::status_translator;
use std::net;
use std::borrow::{Borrow};
use uuid;

#[derive(Debug, Copy, Clone)]
pub enum ConnectionState {
    Establishing{listening: bool, compatible_version: bool, attempts: u32, request_id: Option<u64>},
    Established,
    Closing{request_id: Option<u64>},
    Closed,
}

#[derive(Debug)]
pub struct Connection {
    pub state: ConnectionState,
    pub local_id: u64,
    pub remote_id: u64,
    pub address: net::SocketAddr,
    pub received_packets: u64,
    pub sent_packets: u64,
    pub heartbeat_counter: u64,
    //For echoes.
    pub endpoint_id: uuid::Uuid,
    pub roundtrip_estimator: RoundtripEstimator,
}

const MAX_STATUS_ATTEMPTS: u32 = 10;
const MAX_CONNECTION_ATTEMPTS:u32 = 25; //5000 ms divided by 200 ms per attempt, see spec.

impl Connection {

    pub fn new(address: net::SocketAddr, local_id: u64)->Connection {
        Connection {
            state: ConnectionState::Closed,
            local_id: local_id,
            remote_id: 0,
            address: address,
            sent_packets: 0,
            received_packets: 0,
            heartbeat_counter: 0,
            endpoint_id: uuid::Uuid::new_v4(),
            roundtrip_estimator: RoundtripEstimator::new(5),
        }
    }

    pub fn from_connection_request(address: net::SocketAddr, local_id: u64, remote_id: u64)->Connection {
        let mut conn = Connection::new(address, local_id);
        conn.remote_id = remote_id;
        conn.state = ConnectionState::Established;
        conn
    }

    pub fn establish<H: async::Handler>(&mut self, request_id: Option<u64>, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Closed = self.state {
            self.state = ConnectionState::Establishing{listening: false, compatible_version: false, attempts: 0, request_id: request_id};
            //get things rolling...
            self.send(Packet::StatusRequest(StatusRequest::FastnetQuery), service);
        }
    }

    pub fn send<P: Borrow<Packet>, H: async::Handler>(&mut self, packet: P, service: &mut MioServiceProvider<H>)->bool {
        self.sent_packets += 1;
        service.send(packet, self.address)
    }

    pub fn handle_incoming_packet<H: async::Handler>(&mut self, packet: &Packet, service: &mut MioServiceProvider<H>)->bool {
        self.received_packets += 1; //Always.
        match *packet {
            Packet::StatusResponse(ref resp) => {
                self.handle_status_response(resp, service);
                true
            },
            Packet::Echo{endpoint, uuid} => {
                if endpoint != self.endpoint_id {
                    self.send(packet, service);
                    self.roundtrip_estimator.handle_echo(self.local_id, uuid, service);
                }
                true
            },
            Packet::Heartbeat{counter: c, sent: s, received: r} => {
                true
            },
            Packet::Connected(id) => {
                self.handle_connected(id, service);
                true
            },
            Packet::Aborted(ref message) => {
                self.handle_aborted(message, service);
                true
            },
            _ => false
        }
    }

    fn handle_connected<H: async::Handler>(&mut self, id: u64, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{listening, compatible_version, request_id, ..} = self.state {
            if listening && compatible_version {
                self.remote_id = id;
                //The spec says that heartbeats don't count any packets that happen before full establishment.
                self.sent_packets = 0;
                self.received_packets = 0;
                self.state = ConnectionState::Established;
                service.handler.connected(self.local_id, request_id);
            }
        }
        //Otherwise, we shouldn't be receiving this yet so just drop it.
    }

    fn handle_aborted<H: async::Handler>(&mut self, message: &str, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{listening, compatible_version, request_id, ..} = self.state {
            if listening && compatible_version {
                self.state = ConnectionState::Closed;
                if let Some(id) = request_id {service.handler.request_failed(id, async::Error::ConnectionAborted);}
            }
        }
    }

    fn handle_status_response<H: async::Handler>(&mut self, resp: &StatusResponse, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{mut listening, mut compatible_version, mut attempts, request_id} = self.state {
            match *resp {
                StatusResponse::FastnetResponse(new_listening) if listening == false => {
                    if new_listening == false {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::NotListening);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    listening = true;
                    self.send(Packet::StatusRequest(StatusRequest::VersionQuery), service);
                },
                StatusResponse::VersionResponse(ref v) if compatible_version == false => {
                    if v.eq(status_translator::PROTOCOL_VERSION) == false {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::IncompatibleVersions)}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    compatible_version = true;
                }
                _ => {}
            }
            if listening && compatible_version {
                let id = self.local_id;
                self.send(Packet::Connect(id), service);
            }
            self.state = ConnectionState::Establishing{attempts: 0, listening: listening, compatible_version: compatible_version, request_id: request_id};
        }
    }

    pub fn tick1000<H: async::Handler>(&mut self, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Established = self.state {
            let heartbeat = Packet::Heartbeat{counter: self.heartbeat_counter, sent: self.sent_packets, received: self.received_packets};
            self.heartbeat_counter += 1;
            self.send(heartbeat, service);
        }
    }

    pub fn tick200<H: async::Handler>(&mut self, service: &mut MioServiceProvider<H>) {
        match self.state {
            ConnectionState::Establishing{mut attempts, listening, compatible_version, request_id} => {
                attempts += 1;
                if listening == false {
                    if attempts > MAX_STATUS_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::StatusRequest(StatusRequest::FastnetQuery), self.address);
                }
                else if compatible_version == false {
                    if attempts > MAX_STATUS_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::StatusRequest(StatusRequest::VersionQuery), self.address);
                }
                else {
                    if attempts > MAX_CONNECTION_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::Connect(self.local_id), self.address);
                }
            },
            ConnectionState::Established => {
                self.roundtrip_estimator.tick(self.address, self.endpoint_id, service);
            },
            _ => {},
        }
    }
}
<commit_msg>Fix roundtrip estimation.<commit_after>use super::*;
use super::super::packets::*;
use super::super::async;
use super::super::status_translator;
use std::net;
use std::borrow::{Borrow};
use uuid;

#[derive(Debug, Copy, Clone)]
pub enum ConnectionState {
    Establishing{listening: bool, compatible_version: bool, attempts: u32, request_id: Option<u64>},
    Established,
    Closing{request_id: Option<u64>},
    Closed,
}

#[derive(Debug)]
pub struct Connection {
    pub state: ConnectionState,
    pub local_id: u64,
    pub remote_id: u64,
    pub address: net::SocketAddr,
    pub received_packets: u64,
    pub sent_packets: u64,
    pub heartbeat_counter: u64,
    //For echoes.
    pub endpoint_id: uuid::Uuid,
    pub roundtrip_estimator: RoundtripEstimator,
}

const MAX_STATUS_ATTEMPTS: u32 = 10;
const MAX_CONNECTION_ATTEMPTS:u32 = 25; //5000 ms divided by 200 ms per attempt, see spec.

impl Connection {

    pub fn new(address: net::SocketAddr, local_id: u64)->Connection {
        Connection {
            state: ConnectionState::Closed,
            local_id: local_id,
            remote_id: 0,
            address: address,
            sent_packets: 0,
            received_packets: 0,
            heartbeat_counter: 0,
            endpoint_id: uuid::Uuid::new_v4(),
            roundtrip_estimator: RoundtripEstimator::new(5),
        }
    }

    pub fn from_connection_request(address: net::SocketAddr, local_id: u64, remote_id: u64)->Connection {
        let mut conn = Connection::new(address, local_id);
        conn.remote_id = remote_id;
        conn.state = ConnectionState::Established;
        conn
    }

    pub fn establish<H: async::Handler>(&mut self, request_id: Option<u64>, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Closed = self.state {
            self.state = ConnectionState::Establishing{listening: false, compatible_version: false, attempts: 0, request_id: request_id};
            //get things rolling...
            self.send(Packet::StatusRequest(StatusRequest::FastnetQuery), service);
        }
    }

    pub fn send<P: Borrow<Packet>, H: async::Handler>(&mut self, packet: P, service: &mut MioServiceProvider<H>)->bool {
        self.sent_packets += 1;
        service.send(packet, self.address)
    }

    pub fn handle_incoming_packet<H: async::Handler>(&mut self, packet: &Packet, service: &mut MioServiceProvider<H>)->bool {
        self.received_packets += 1; //Always.
        match *packet {
            Packet::StatusResponse(ref resp) => {
                self.handle_status_response(resp, service);
                true
            },
            Packet::Echo{endpoint, uuid} => {
                if endpoint != self.endpoint_id {
                    self.send(packet, service);
                }
                else {
                    self.roundtrip_estimator.handle_echo(self.local_id, uuid, service);
                }
                true
            },
            Packet::Heartbeat{counter: c, sent: s, received: r} => {
                true
            },
            Packet::Connected(id) => {
                self.handle_connected(id, service);
                true
            },
            Packet::Aborted(ref message) => {
                self.handle_aborted(message, service);
                true
            },
            _ => false
        }
    }

    fn handle_connected<H: async::Handler>(&mut self, id: u64, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{listening, compatible_version, request_id, ..} = self.state {
            if listening && compatible_version {
                self.remote_id = id;
                //The spec says that heartbeats don't count any packets that happen before full establishment.
                self.sent_packets = 0;
                self.received_packets = 0;
                self.state = ConnectionState::Established;
                service.handler.connected(self.local_id, request_id);
            }
        }
        //Otherwise, we shouldn't be receiving this yet so just drop it.
    }

    fn handle_aborted<H: async::Handler>(&mut self, message: &str, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{listening, compatible_version, request_id, ..} = self.state {
            if listening && compatible_version {
                self.state = ConnectionState::Closed;
                if let Some(id) = request_id {service.handler.request_failed(id, async::Error::ConnectionAborted);}
            }
        }
    }

    fn handle_status_response<H: async::Handler>(&mut self, resp: &StatusResponse, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Establishing{mut listening, mut compatible_version, mut attempts, request_id} = self.state {
            match *resp {
                StatusResponse::FastnetResponse(new_listening) if listening == false => {
                    if new_listening == false {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::NotListening);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    listening = true;
                    self.send(Packet::StatusRequest(StatusRequest::VersionQuery), service);
                },
                StatusResponse::VersionResponse(ref v) if compatible_version == false => {
                    if v.eq(status_translator::PROTOCOL_VERSION) == false {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::IncompatibleVersions)}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    compatible_version = true;
                }
                _ => {}
            }
            if listening && compatible_version {
                let id = self.local_id;
                self.send(Packet::Connect(id), service);
            }
            self.state = ConnectionState::Establishing{attempts: 0, listening: listening, compatible_version: compatible_version, request_id: request_id};
        }
    }

    pub fn tick1000<H: async::Handler>(&mut self, service: &mut MioServiceProvider<H>) {
        if let ConnectionState::Established = self.state {
            let heartbeat = Packet::Heartbeat{counter: self.heartbeat_counter, sent: self.sent_packets, received: self.received_packets};
            self.heartbeat_counter += 1;
            self.send(heartbeat, service);
        }
    }

    pub fn tick200<H: async::Handler>(&mut self, service: &mut MioServiceProvider<H>) {
        match self.state {
            ConnectionState::Establishing{mut attempts, listening, compatible_version, request_id} => {
                attempts += 1;
                if listening == false {
                    if attempts > MAX_STATUS_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::StatusRequest(StatusRequest::FastnetQuery), self.address);
                }
                else if compatible_version == false {
                    if attempts > MAX_STATUS_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::StatusRequest(StatusRequest::VersionQuery), self.address);
                }
                else {
                    if attempts > MAX_CONNECTION_ATTEMPTS {
                        if let Some(id) = request_id {service.handler.request_failed(id, async::Error::TimedOut);}
                        self.state = ConnectionState::Closed;
                        return;
                    }
                    service.send(Packet::Connect(self.local_id), self.address);
                }
            },
            ConnectionState::Established => {
                self.roundtrip_estimator.tick(self.address, self.endpoint_id, service);
            },
            _ => {},
        }
    }
}
<|endoftext|>


<commit_before>//! Solve Hermitian (or real symmetric) linear problems and invert Hermitian
//! (or real symmetric) matrices
//!
//! **Note that only the upper triangular portion of the matrix is used.**
//!
//! # Examples
//!
//! Solve `A * x = b`, where `A` is a Hermitian (or real symmetric) matrix:
//!
//! ```
//! #[macro_use]
//! extern crate ndarray;
//! extern crate ndarray_linalg;
//!
//! use ndarray::prelude::*;
//! use ndarray_linalg::SolveH;
//! # fn main() {
//!
//! let a: Array2<f64> = array![
//!     [3., 2., -1.],
//!     [2., -2., 4.],
//!     [-1., 4., 5.]
//! ];
//! let b: Array1<f64> = array![11., -12., 1.];
//! let x = a.solveh_into(b).unwrap();
//! assert!(x.all_close(&array![1., 3., -2.], 1e-9));
//!
//! # }
//! ```
//!
//! If you are solving multiple systems of linear equations with the same
//! Hermitian or real symmetric coefficient matrix `A`, it's faster to compute
//! the factorization once at the beginning than solving directly using `A`:
//!
//! ```
//! # extern crate ndarray;
//! # extern crate ndarray_linalg;
//! use ndarray::prelude::*;
//! use ndarray_linalg::*;
//! # fn main() {
//!
//! let a: Array2<f64> = random((3, 3));
//! let f = a.factorizeh_into().unwrap(); // Factorize A (A is consumed)
//! for _ in 0..10 {
//!     let b: Array1<f64> = random(3);
//!     let x = f.solveh_into(b).unwrap(); // Solve A * x = b using the factorization
//! }
//!
//! # }
//! ```

use ndarray::*;

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

pub use lapack_traits::{Pivot, UPLO};

/// An interface for solving systems of Hermitian (or real symmetric) linear equations.
///
/// If you plan to solve many equations with the same Hermitian (or real
/// symmetric) coefficient matrix `A` but different `b` vectors, it's faster to
/// factor the `A` matrix once using the `FactorizeH` trait, and then solve
/// using the `BKFactorized` struct.
pub trait SolveH<A: Scalar> {
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result.
    fn solveh<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solveh_inplace(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result.
    fn solveh_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solveh_inplace(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result. The value of `x` is also assigned to the
    /// argument.
    fn solveh_inplace<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
}

/// Represents the Bunch–Kaufman factorization of a Hermitian (or real
/// symmetric) matrix as `A = P * U * D * U^H * P^T`.
pub struct BKFactorized<S: Data> {
    pub a: ArrayBase<S, Ix2>,
    pub ipiv: Pivot,
}

impl<A, S> SolveH<A> for BKFactorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solveh_inplace<'a, Sb>(&self, rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solveh(
                self.a.square_layout()?,
                UPLO::Upper,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
}

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

/// An interface for computing the Bunch–Kaufman factorization of Hermitian (or
/// real symmetric) matrix refs.
pub trait FactorizeH<S: Data> {
    /// Computes the Bunch–Kaufman factorization of a Hermitian (or real
    /// symmetric) matrix.
    fn factorizeh(&self) -> Result<BKFactorized<S>>;
}

/// An interface for computing the Bunch–Kaufman factorization of Hermitian (or
/// real symmetric) matrices.
pub trait FactorizeHInto<S: Data> {
    /// Computes the Bunch–Kaufman factorization of a Hermitian (or real
    /// symmetric) matrix.
    fn factorizeh_into(self) -> Result<BKFactorized<S>>;
}

impl<A, S> FactorizeHInto<S> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    fn factorizeh_into(mut self) -> Result<BKFactorized<S>> {
        let ipiv = unsafe { A::bk(self.layout()?, UPLO::Upper, self.as_allocated_mut()?)? };
        Ok(BKFactorized {
            a: self,
            ipiv: ipiv,
        })
    }
}

impl<A, Si> FactorizeH<OwnedRepr<A>> for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    fn factorizeh(&self) -> Result<BKFactorized<OwnedRepr<A>>> {
        let mut a: Array2<A> = replicate(self);
        let ipiv = unsafe { A::bk(a.layout()?, UPLO::Upper, a.as_allocated_mut()?)? };
        Ok(BKFactorized { a: a, ipiv: ipiv })
    }
}

/// An interface for inverting Hermitian (or real symmetric) matrix refs.
pub trait InverseH {
    type Output;
    /// Computes the inverse of the Hermitian (or real symmetric) matrix.
    ///
    /// **Warning: The inverse is stored only in the upper triangular portion
    /// of the result matrix!** If you want the lower triangular portion to be
    /// correct, you must fill it in according to the results in the upper
    /// triangular portion.
    fn invh(&self) -> Result<Self::Output>;
}

/// An interface for inverting Hermitian (or real symmetric) matrices.
pub trait InverseHInto {
    type Output;
    /// Computes the inverse of the Hermitian (or real symmetric) matrix.
    ///
    /// **Warning: The inverse is stored only in the upper triangular portion
    /// of the result matrix!** If you want the lower triangular portion to be
    /// correct, you must fill it in according to the results in the upper
    /// triangular portion.
    fn invh_into(self) -> Result<Self::Output>;
}

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

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

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

    fn invh(&self) -> Result<Self::Output> {
        let f = BKFactorized {
            a: replicate(&self.a),
            ipiv: self.ipiv.clone(),
        };
        f.invh_into()
    }
}

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

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

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

    fn invh(&self) -> Result<Self::Output> {
        let f = self.factorizeh()?;
        f.invh_into()
    }
}
<commit_msg>Change `factorizeh*` to require square input<commit_after>//! Solve Hermitian (or real symmetric) linear problems and invert Hermitian
//! (or real symmetric) matrices
//!
//! **Note that only the upper triangular portion of the matrix is used.**
//!
//! # Examples
//!
//! Solve `A * x = b`, where `A` is a Hermitian (or real symmetric) matrix:
//!
//! ```
//! #[macro_use]
//! extern crate ndarray;
//! extern crate ndarray_linalg;
//!
//! use ndarray::prelude::*;
//! use ndarray_linalg::SolveH;
//! # fn main() {
//!
//! let a: Array2<f64> = array![
//!     [3., 2., -1.],
//!     [2., -2., 4.],
//!     [-1., 4., 5.]
//! ];
//! let b: Array1<f64> = array![11., -12., 1.];
//! let x = a.solveh_into(b).unwrap();
//! assert!(x.all_close(&array![1., 3., -2.], 1e-9));
//!
//! # }
//! ```
//!
//! If you are solving multiple systems of linear equations with the same
//! Hermitian or real symmetric coefficient matrix `A`, it's faster to compute
//! the factorization once at the beginning than solving directly using `A`:
//!
//! ```
//! # extern crate ndarray;
//! # extern crate ndarray_linalg;
//! use ndarray::prelude::*;
//! use ndarray_linalg::*;
//! # fn main() {
//!
//! let a: Array2<f64> = random((3, 3));
//! let f = a.factorizeh_into().unwrap(); // Factorize A (A is consumed)
//! for _ in 0..10 {
//!     let b: Array1<f64> = random(3);
//!     let x = f.solveh_into(b).unwrap(); // Solve A * x = b using the factorization
//! }
//!
//! # }
//! ```

use ndarray::*;

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

pub use lapack_traits::{Pivot, UPLO};

/// An interface for solving systems of Hermitian (or real symmetric) linear equations.
///
/// If you plan to solve many equations with the same Hermitian (or real
/// symmetric) coefficient matrix `A` but different `b` vectors, it's faster to
/// factor the `A` matrix once using the `FactorizeH` trait, and then solve
/// using the `BKFactorized` struct.
pub trait SolveH<A: Scalar> {
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result.
    fn solveh<S: Data<Elem = A>>(&self, b: &ArrayBase<S, Ix1>) -> Result<Array1<A>> {
        let mut b = replicate(b);
        self.solveh_inplace(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result.
    fn solveh_into<S: DataMut<Elem = A>>(&self, mut b: ArrayBase<S, Ix1>) -> Result<ArrayBase<S, Ix1>> {
        self.solveh_inplace(&mut b)?;
        Ok(b)
    }
    /// Solves a system of linear equations `A * x = b` with Hermitian (or real
    /// symmetric) matrix `A`, where `A` is `self`, `b` is the argument, and
    /// `x` is the successful result. The value of `x` is also assigned to the
    /// argument.
    fn solveh_inplace<'a, S: DataMut<Elem = A>>(&self, &'a mut ArrayBase<S, Ix1>) -> Result<&'a mut ArrayBase<S, Ix1>>;
}

/// Represents the Bunch–Kaufman factorization of a Hermitian (or real
/// symmetric) matrix as `A = P * U * D * U^H * P^T`.
pub struct BKFactorized<S: Data> {
    pub a: ArrayBase<S, Ix2>,
    pub ipiv: Pivot,
}

impl<A, S> SolveH<A> for BKFactorized<S>
where
    A: Scalar,
    S: Data<Elem = A>,
{
    fn solveh_inplace<'a, Sb>(&self, rhs: &'a mut ArrayBase<Sb, Ix1>) -> Result<&'a mut ArrayBase<Sb, Ix1>>
    where
        Sb: DataMut<Elem = A>,
    {
        unsafe {
            A::solveh(
                self.a.square_layout()?,
                UPLO::Upper,
                self.a.as_allocated()?,
                &self.ipiv,
                rhs.as_slice_mut().unwrap(),
            )?
        };
        Ok(rhs)
    }
}

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

/// An interface for computing the Bunch–Kaufman factorization of Hermitian (or
/// real symmetric) matrix refs.
pub trait FactorizeH<S: Data> {
    /// Computes the Bunch–Kaufman factorization of a Hermitian (or real
    /// symmetric) matrix.
    fn factorizeh(&self) -> Result<BKFactorized<S>>;
}

/// An interface for computing the Bunch–Kaufman factorization of Hermitian (or
/// real symmetric) matrices.
pub trait FactorizeHInto<S: Data> {
    /// Computes the Bunch–Kaufman factorization of a Hermitian (or real
    /// symmetric) matrix.
    fn factorizeh_into(self) -> Result<BKFactorized<S>>;
}

impl<A, S> FactorizeHInto<S> for ArrayBase<S, Ix2>
where
    A: Scalar,
    S: DataMut<Elem = A>,
{
    fn factorizeh_into(mut self) -> Result<BKFactorized<S>> {
        let ipiv = unsafe { A::bk(self.square_layout()?, UPLO::Upper, self.as_allocated_mut()?)? };
        Ok(BKFactorized {
            a: self,
            ipiv: ipiv,
        })
    }
}

impl<A, Si> FactorizeH<OwnedRepr<A>> for ArrayBase<Si, Ix2>
where
    A: Scalar,
    Si: Data<Elem = A>,
{
    fn factorizeh(&self) -> Result<BKFactorized<OwnedRepr<A>>> {
        let mut a: Array2<A> = replicate(self);
        let ipiv = unsafe { A::bk(a.square_layout()?, UPLO::Upper, a.as_allocated_mut()?)? };
        Ok(BKFactorized { a: a, ipiv: ipiv })
    }
}

/// An interface for inverting Hermitian (or real symmetric) matrix refs.
pub trait InverseH {
    type Output;
    /// Computes the inverse of the Hermitian (or real symmetric) matrix.
    ///
    /// **Warning: The inverse is stored only in the upper triangular portion
    /// of the result matrix!** If you want the lower triangular portion to be
    /// correct, you must fill it in according to the results in the upper
    /// triangular portion.
    fn invh(&self) -> Result<Self::Output>;
}

/// An interface for inverting Hermitian (or real symmetric) matrices.
pub trait InverseHInto {
    type Output;
    /// Computes the inverse of the Hermitian (or real symmetric) matrix.
    ///
    /// **Warning: The inverse is stored only in the upper triangular portion
    /// of the result matrix!** If you want the lower triangular portion to be
    /// correct, you must fill it in according to the results in the upper
    /// triangular portion.
    fn invh_into(self) -> Result<Self::Output>;
}

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

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

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

    fn invh(&self) -> Result<Self::Output> {
        let f = BKFactorized {
            a: replicate(&self.a),
            ipiv: self.ipiv.clone(),
        };
        f.invh_into()
    }
}

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

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

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

    fn invh(&self) -> Result<Self::Output> {
        let f = self.factorizeh()?;
        f.invh_into()
    }
}
<|endoftext|>


<commit_before>// ignore-windows: Concurrency on Windows is not supported yet.
// compile-flags: -Zmiri-ignore-leaks -Zmiri-disable-stacked-borrows

// The following tests check whether our weak memory emulation produces
// any inconsistent execution outcomes
//
// Due to the random nature of choosing valid stores, it is always
// possible that our tests spuriously succeeds: even though our weak
// memory emulation code has incorrectly identified a store in
// modification order as being valid, it may be never chosen by
// the RNG and never observed in our tests.
//
// To mitigate this, each test is ran enough times such that the chance
// of spurious success is very low. These tests never supriously fail.

// Test cases and their consistent outcomes are from
// http://svr-pes20-cppmem.cl.cam.ac.uk/cppmem/
// Based on
// M. Batty, S. Owens, S. Sarkar, P. Sewell and T. Weber,
// "Mathematizing C++ concurrency", ACM SIGPLAN Notices, vol. 46, no. 1, pp. 55-66, 2011.
// Available: https://ss265.host.cs.st-andrews.ac.uk/papers/n3132.pdf.

use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::*;
use std::thread::{spawn, yield_now};

#[derive(Copy, Clone)]
struct EvilSend<T>(pub T);

unsafe impl<T> Send for EvilSend<T> {}
unsafe impl<T> Sync for EvilSend<T> {}

// We can't create static items because we need to run each test
// multiple times
fn static_atomic(val: usize) -> &'static AtomicUsize {
    let ret = Box::leak(Box::new(AtomicUsize::new(val)));
    // A workaround to put the initialisation value in the store buffer
    ret.store(val, Relaxed);
    ret
}

// Spins and yields until until acquires a pre-determined value
fn acquires_value(loc: &AtomicUsize, val: usize) -> usize {
    while loc.load(Acquire) != val {
        yield_now();
    }
    val
}

fn test_corr() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Relaxed);
        x.store(2, Relaxed);
    });

    let j2 = spawn(move || {
        let r2 = x.load(Relaxed); // -------------------------------------+
        y.store(1, Release); // ---------------------+                    |
        r2 //                                        |                    |
    }); //                                           |                    |
    //                                               |synchronizes-with   |happens-before
    let j3 = spawn(move || { //                      |                    |
        acquires_value(&y, 1); // <------------------+                    |
        x.load(Relaxed) // <----------------------------------------------+
        // The two reads on x are ordered by hb, so they cannot observe values
        // differently from the modification order. If the first read observed
        // 2, then the second read must observe 2 as well.
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();
    let r3 = j3.join().unwrap();
    if r2 == 2 {
        assert_eq!(r3, 2);
    }
}

fn test_wrc() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Release); // ---------------------+---------------------+
    }); //                                           |                     |
    //                                               |synchronizes-with    |
    let j2 = spawn(move || { //                      |                     |
        acquires_value(&x, 1); // <------------------+                     |
        y.store(1, Release); // ---------------------+                     |happens-before
    }); //                                           |                     |
    //                                               |synchronizes-with    |
    let j3 = spawn(move || { //                      |                     |
        acquires_value(&y, 1); // <------------------+                     |
        x.load(Relaxed) // <-----------------------------------------------+
    });

    j1.join().unwrap();
    j2.join().unwrap();
    let r3 = j3.join().unwrap();

    assert_eq!(r3, 1);
}

fn test_message_passing() {
    let mut var = 0u32;
    let ptr = &mut var as *mut u32;
    let x = EvilSend(ptr);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        unsafe { *x.0 = 1 }; // -----------------------------------------+
        y.store(1, Release); // ---------------------+                   |
    }); //                                           |                   |
    //                                               |synchronizes-with  | happens-before
    let j2 = spawn(move || { //                      |                   |
        acquires_value(&y, 1); // <------------------+                   |
        unsafe { *x.0 } // <---------------------------------------------+
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    assert_eq!(r2, 1);
}

// LB+acq_rel+acq_rel
fn test_load_buffering_acq_rel() {
    let x = static_atomic(0);
    let y = static_atomic(0);
    let j1 = spawn(move || {
        let r1 = x.load(Acquire);
        y.store(1, Release);
        r1
    });

    let j2 = spawn(move || {
        let r2 = y.load(Acquire);
        x.store(1, Release);
        r2
    });

    let r1 = j1.join().unwrap();
    let r2 = j2.join().unwrap();

    // 3 consistent outcomes: (0,0), (0,1), (1,0)
    assert_ne!((r1, r2), (1, 1));
}

fn test_mixed_access() {
    /*
    int main() {
      atomic_int x = 0;
      {{{
        x.store(1, mo_relaxed);
      }}}

      x.store(2, mo_relaxed);

      {{{
        r1 = x.load(mo_relaxed);
      }}}

      return 0;
    }
        */
    let x = static_atomic(0);

    spawn(move || {
        x.store(1, Relaxed);
    })
    .join()
    .unwrap();

    x.store(2, Relaxed);

    let r2 = spawn(move || x.load(Relaxed)).join().unwrap();

    assert_eq!(r2, 2);
}

// The following two tests are taken from Repairing Sequential Consistency in C/C++11
// by Lahav et al.
// https://plv.mpi-sws.org/scfix/paper.pdf

// Test case SB
fn test_sc_store_buffering() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, SeqCst);
        y.load(SeqCst)
    });

    let j2 = spawn(move || {
        y.store(1, SeqCst);
        x.load(SeqCst)
    });

    let a = j1.join().unwrap();
    let b = j2.join().unwrap();

    assert_ne!((a, b), (0, 0));
}

pub fn main() {
    // TODO: does this make chances of spurious success
    // "sufficiently low"? This also takes a long time to run,
    // prehaps each function should be its own test case so they
    // can be run in parallel
    for _ in 0..500 {
        test_mixed_access();
        test_load_buffering_acq_rel();
        test_message_passing();
        test_wrc();
        test_corr();
        test_sc_store_buffering();
    }
}
<commit_msg>Reduce the number of runs in consistency tests<commit_after>// ignore-windows: Concurrency on Windows is not supported yet.
// compile-flags: -Zmiri-ignore-leaks -Zmiri-disable-stacked-borrows

// The following tests check whether our weak memory emulation produces
// any inconsistent execution outcomes
//
// Due to the random nature of choosing valid stores, it is always
// possible that our tests spuriously succeeds: even though our weak
// memory emulation code has incorrectly identified a store in
// modification order as being valid, it may be never chosen by
// the RNG and never observed in our tests.
//
// To mitigate this, each test is ran enough times such that the chance
// of spurious success is very low. These tests never supriously fail.

// Test cases and their consistent outcomes are from
// http://svr-pes20-cppmem.cl.cam.ac.uk/cppmem/
// Based on
// M. Batty, S. Owens, S. Sarkar, P. Sewell and T. Weber,
// "Mathematizing C++ concurrency", ACM SIGPLAN Notices, vol. 46, no. 1, pp. 55-66, 2011.
// Available: https://ss265.host.cs.st-andrews.ac.uk/papers/n3132.pdf.

use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::*;
use std::thread::{spawn, yield_now};

#[derive(Copy, Clone)]
struct EvilSend<T>(pub T);

unsafe impl<T> Send for EvilSend<T> {}
unsafe impl<T> Sync for EvilSend<T> {}

// We can't create static items because we need to run each test
// multiple times
fn static_atomic(val: usize) -> &'static AtomicUsize {
    let ret = Box::leak(Box::new(AtomicUsize::new(val)));
    // A workaround to put the initialisation value in the store buffer
    ret.store(val, Relaxed);
    ret
}

// Spins and yields until until acquires a pre-determined value
fn acquires_value(loc: &AtomicUsize, val: usize) -> usize {
    while loc.load(Acquire) != val {
        yield_now();
    }
    val
}

fn test_corr() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Relaxed);
        x.store(2, Relaxed);
    });

    let j2 = spawn(move || {
        let r2 = x.load(Relaxed); // -------------------------------------+
        y.store(1, Release); // ---------------------+                    |
        r2 //                                        |                    |
    }); //                                           |                    |
    //                                               |synchronizes-with   |happens-before
    let j3 = spawn(move || { //                      |                    |
        acquires_value(&y, 1); // <------------------+                    |
        x.load(Relaxed) // <----------------------------------------------+
        // The two reads on x are ordered by hb, so they cannot observe values
        // differently from the modification order. If the first read observed
        // 2, then the second read must observe 2 as well.
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();
    let r3 = j3.join().unwrap();
    if r2 == 2 {
        assert_eq!(r3, 2);
    }
}

fn test_wrc() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Release); // ---------------------+---------------------+
    }); //                                           |                     |
    //                                               |synchronizes-with    |
    let j2 = spawn(move || { //                      |                     |
        acquires_value(&x, 1); // <------------------+                     |
        y.store(1, Release); // ---------------------+                     |happens-before
    }); //                                           |                     |
    //                                               |synchronizes-with    |
    let j3 = spawn(move || { //                      |                     |
        acquires_value(&y, 1); // <------------------+                     |
        x.load(Relaxed) // <-----------------------------------------------+
    });

    j1.join().unwrap();
    j2.join().unwrap();
    let r3 = j3.join().unwrap();

    assert_eq!(r3, 1);
}

fn test_message_passing() {
    let mut var = 0u32;
    let ptr = &mut var as *mut u32;
    let x = EvilSend(ptr);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        unsafe { *x.0 = 1 }; // -----------------------------------------+
        y.store(1, Release); // ---------------------+                   |
    }); //                                           |                   |
    //                                               |synchronizes-with  | happens-before
    let j2 = spawn(move || { //                      |                   |
        acquires_value(&y, 1); // <------------------+                   |
        unsafe { *x.0 } // <---------------------------------------------+
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    assert_eq!(r2, 1);
}

// LB+acq_rel+acq_rel
fn test_load_buffering_acq_rel() {
    let x = static_atomic(0);
    let y = static_atomic(0);
    let j1 = spawn(move || {
        let r1 = x.load(Acquire);
        y.store(1, Release);
        r1
    });

    let j2 = spawn(move || {
        let r2 = y.load(Acquire);
        x.store(1, Release);
        r2
    });

    let r1 = j1.join().unwrap();
    let r2 = j2.join().unwrap();

    // 3 consistent outcomes: (0,0), (0,1), (1,0)
    assert_ne!((r1, r2), (1, 1));
}

fn test_mixed_access() {
    /*
    int main() {
      atomic_int x = 0;
      {{{
        x.store(1, mo_relaxed);
      }}}

      x.store(2, mo_relaxed);

      {{{
        r1 = x.load(mo_relaxed);
      }}}

      return 0;
    }
        */
    let x = static_atomic(0);

    spawn(move || {
        x.store(1, Relaxed);
    })
    .join()
    .unwrap();

    x.store(2, Relaxed);

    let r2 = spawn(move || x.load(Relaxed)).join().unwrap();

    assert_eq!(r2, 2);
}

// The following two tests are taken from Repairing Sequential Consistency in C/C++11
// by Lahav et al.
// https://plv.mpi-sws.org/scfix/paper.pdf

// Test case SB
fn test_sc_store_buffering() {
    let x = static_atomic(0);
    let y = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, SeqCst);
        y.load(SeqCst)
    });

    let j2 = spawn(move || {
        y.store(1, SeqCst);
        x.load(SeqCst)
    });

    let a = j1.join().unwrap();
    let b = j2.join().unwrap();

    assert_ne!((a, b), (0, 0));
}

pub fn main() {
    for _ in 0..100 {
        test_mixed_access();
        test_load_buffering_acq_rel();
        test_message_passing();
        test_wrc();
        test_corr();
        test_sc_store_buffering();
    }
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use crate::image::manifest::Manifest;
use crate::image::manifest_def::ManifestDef;
use crate::util::parse_int::ParseInt;
use anyhow::{ensure, Result};
use std::convert::TryInto;
use std::fs::File;
use std::io::{Read, Write};
use std::mem::size_of;
use std::path::{Path, PathBuf};
use thiserror::Error;

use zerocopy::LayoutVerified;

#[derive(Debug, Error)]
pub enum ImageError {
    #[error("Incomplete read: expected to read {0} bytes but read {1} bytes")]
    IncompleteRead(usize, usize),
    #[error("Failed to parse image manifest.")]
    Parse,
}

#[repr(C)]
#[derive(Debug)]
pub struct ImageData {
    pub bytes: [u8; Image::MAX_SIZE],
    _align: [Manifest; 0],
}

impl Default for ImageData {
    fn default() -> Self {
        ImageData {
            bytes: [0xFF; Image::MAX_SIZE],
            _align: [],
        }
    }
}

#[derive(Debug, Default)]
pub struct Image {
    data: Box<ImageData>,
    size: usize,
}

#[derive(Debug)]
pub enum ImageChunk {
    Concat(PathBuf),
    Offset(PathBuf, usize),
}

#[derive(Debug, Default)]
pub struct ImageAssembler {
    pub size: usize,
    pub mirrored: bool,
    pub chunks: Vec<ImageChunk>,
}

impl Image {
    const MAX_SIZE: usize = 512 * 1024;
    /// Creates an `Image` from a given input binary.
    pub fn from_reader(mut r: impl Read) -> Result<Self> {
        let mut image = Image::default();
        image.size = r.read(&mut image.data.bytes)?;
        Ok(image)
    }

    pub fn to_writer(&self, w: &mut impl Write) -> Result<()> {
        w.write_all(&self.data.bytes[..self.size])?;
        Ok(())
    }

    pub fn read_from_file(path: &Path) -> Result<Image> {
        let file = File::open(path)?;
        Self::from_reader(file)
    }

    /// Write out the `Image` to a file at the given `path`.
    pub fn write_to_file(self, path: &Path) -> Result<()> {
        let mut file = File::create(path)?;
        self.to_writer(&mut file)
    }

    /// Overwrites all fields in the image's manifest that are defined in `other`.
    pub fn overwrite_manifest(&mut self, other: ManifestDef) -> Result<()> {
        let manifest_slice = &mut self.data.bytes[0..size_of::<Manifest>()];
        let manifest_layout: LayoutVerified<&mut [u8], Manifest> =
            LayoutVerified::new(&mut *manifest_slice).ok_or(ImageError::Parse)?;
        let manifest: &mut Manifest = manifest_layout.into_mut();
        let mut manifest_def: ManifestDef = (&*manifest).try_into()?;
        manifest_def.overwrite_fields(other);
        *manifest = manifest_def.try_into()?;
        Ok(())
    }
}

impl ImageAssembler {
    /// Creates an `ImageAssembler` with a given `size` and mirroring parameters.
    pub fn with_params(size: usize, mirrored: bool) -> Self {
        ImageAssembler {
            size,
            mirrored,
            ..Default::default()
        }
    }

    /// Creates an `ImageAssembler` with default parameters for OpenTitan: a 1MiB image which is mirrored.
    pub fn new() -> Self {
        Self::with_params(0x100000, true)
    }

    /// Parse a list of strings into chunks to be assembled.
    /// Each string may be a filename or a filename@offset describing where in the assembled image the contents of the file should appear.
    /// The offset is an integer expressed in any of the bases accepted by [`ParseInt`].
    pub fn parse(&mut self, chunks: &[impl AsRef<str>]) -> Result<()> {
        for chunk in chunks {
            if let Some((file, offset)) = chunk.as_ref().split_once('@') {
                self.chunks.push(ImageChunk::Offset(
                    PathBuf::from(file),
                    usize::from_str(offset)?,
                ));
            } else {
                self.chunks
                    .push(ImageChunk::Concat(PathBuf::from(chunk.as_ref())));
            }
        }
        Ok(())
    }

    // Read a file into a buffer.  Ensure the entire file is read.
    fn read(path: &Path, buf: &mut [u8]) -> Result<usize> {
        let mut file = File::open(path)?;
        let len = file.metadata()?.len() as usize;
        let n = file.read(buf)?;
        ensure!(len == n, ImageError::IncompleteRead(len, n));
        Ok(n)
    }

    /// Assemble the image according to the parameters and parsed chunk specifications.
    pub fn assemble(&self) -> Result<Vec<u8>> {
        let size = if self.mirrored {
            self.size / 2
        } else {
            self.size
        };
        let mut image = vec![0xff; size];
        let mut pos = 0;
        for chunk in &self.chunks {
            match chunk {
                ImageChunk::Concat(path) => {
                    let n = Self::read(path, &mut image[pos..])?;
                    pos += n;
                }
                ImageChunk::Offset(path, offset) => {
                    let n = Self::read(path, &mut image[*offset..])?;
                    pos = offset + n;
                }
            }
        }
        if self.mirrored {
            image.extend_from_within(..size);
        }
        Ok(image)
    }
}

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

    #[test]
    fn test_assemble_concat() -> Result<()> {
        // Test image assembly by concatenation.
        let mut image = ImageAssembler::with_params(16, false);
        image.parse(&[
            testdata!("hello.txt").to_str().unwrap(),
            testdata!("world.txt").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"HelloWorld\xff\xff\xff\xff\xff\xff");
        Ok(())
    }

    #[test]
    fn test_assemble_offset() -> Result<()> {
        // Test image assembly by explicit offsets.
        let mut image = ImageAssembler::with_params(16, false);
        image.parse(&[
            testdata!("hello.txt@0").to_str().unwrap(),
            testdata!("world.txt@0x8").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"Hello\xff\xff\xffWorld\xff\xff\xff");
        Ok(())
    }

    #[test]
    fn test_assemble_mirrored() -> Result<()> {
        // Test image assembly with mirroring.
        let mut image = ImageAssembler::with_params(20, true);
        image.parse(&[
            testdata!("hello.txt").to_str().unwrap(),
            testdata!("world.txt").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"HelloWorldHelloWorld");
        Ok(())
    }

    #[test]
    fn test_assemble_mirrored_offset_error() -> Result<()> {
        // Test image assembly where one of the source files isn't read completely.
        let mut image = ImageAssembler::with_params(16, true);
        image.parse(&[
            testdata!("hello.txt@0").to_str().unwrap(),
            testdata!("world.txt@0x5").to_str().unwrap(),
        ])?;
        let err = image.assemble().unwrap_err();
        assert_eq!(
            err.to_string(),
            "Incomplete read: expected to read 5 bytes but read 3 bytes"
        );
        Ok(())
    }

    #[test]
    fn test_load_image() {
        // Read and write back image.
        let image = Image::read_from_file(&testdata!("test_image.bin")).unwrap();
        image
            .write_to_file(&testdata!("test_image_out.bin"))
            .unwrap();

        // Ensure the result is identical to the original.
        let (mut orig_bytes, mut res_bytes) = (Vec::<u8>::new(), Vec::<u8>::new());
        File::open(&testdata!("test_image.bin"))
            .unwrap()
            .read_to_end(&mut orig_bytes)
            .unwrap();
        File::open(&testdata!("test_image_out.bin"))
            .unwrap()
            .read_to_end(&mut res_bytes)
            .unwrap();
        assert_eq!(orig_bytes, res_bytes);
    }
}
<commit_msg>[opentitanlib] Add doc comments for Image routines<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use crate::image::manifest::Manifest;
use crate::image::manifest_def::ManifestDef;
use crate::util::parse_int::ParseInt;
use anyhow::{ensure, Result};
use std::convert::TryInto;
use std::fs::File;
use std::io::{Read, Write};
use std::mem::size_of;
use std::path::{Path, PathBuf};
use thiserror::Error;

use zerocopy::LayoutVerified;

#[derive(Debug, Error)]
pub enum ImageError {
    #[error("Incomplete read: expected to read {0} bytes but read {1} bytes")]
    IncompleteRead(usize, usize),
    #[error("Failed to parse image manifest.")]
    Parse,
}

/// A buffer with the same alignment as `Manifest` for storing image data.
#[repr(C)]
#[derive(Debug)]
pub struct ImageData {
    pub bytes: [u8; Image::MAX_SIZE],
    _align: [Manifest; 0],
}

impl Default for ImageData {
    fn default() -> Self {
        ImageData {
            bytes: [0xFF; Image::MAX_SIZE],
            _align: [],
        }
    }
}

#[derive(Debug, Default)]
pub struct Image {
    data: Box<ImageData>,
    size: usize,
}

#[derive(Debug)]
pub enum ImageChunk {
    Concat(PathBuf),
    Offset(PathBuf, usize),
}

#[derive(Debug, Default)]
pub struct ImageAssembler {
    pub size: usize,
    pub mirrored: bool,
    pub chunks: Vec<ImageChunk>,
}

impl Image {
    const MAX_SIZE: usize = 512 * 1024;
    /// Reads in an `Image`.
    pub fn from_reader(mut r: impl Read) -> Result<Self> {
        let mut image = Image::default();
        image.size = r.read(&mut image.data.bytes)?;
        Ok(image)
    }

    /// Writes out the `Image`.
    pub fn to_writer(&self, w: &mut impl Write) -> Result<()> {
        w.write_all(&self.data.bytes[..self.size])?;
        Ok(())
    }

    /// Creates an `Image` from a given input binary.
    pub fn read_from_file(path: &Path) -> Result<Image> {
        let file = File::open(path)?;
        Self::from_reader(file)
    }

    /// Write out the `Image` to a file at the given `path`.
    pub fn write_to_file(self, path: &Path) -> Result<()> {
        let mut file = File::create(path)?;
        self.to_writer(&mut file)
    }

    /// Overwrites all fields in the image's manifest that are defined in `other`.
    pub fn overwrite_manifest(&mut self, other: ManifestDef) -> Result<()> {
        let manifest_slice = &mut self.data.bytes[0..size_of::<Manifest>()];
        let manifest_layout: LayoutVerified<&mut [u8], Manifest> =
            LayoutVerified::new(&mut *manifest_slice).ok_or(ImageError::Parse)?;
        let manifest: &mut Manifest = manifest_layout.into_mut();
        let mut manifest_def: ManifestDef = (&*manifest).try_into()?;
        manifest_def.overwrite_fields(other);
        *manifest = manifest_def.try_into()?;
        Ok(())
    }
}

impl ImageAssembler {
    /// Creates an `ImageAssembler` with a given `size` and mirroring parameters.
    pub fn with_params(size: usize, mirrored: bool) -> Self {
        ImageAssembler {
            size,
            mirrored,
            ..Default::default()
        }
    }

    /// Creates an `ImageAssembler` with default parameters for OpenTitan: a 1MiB image which is mirrored.
    pub fn new() -> Self {
        Self::with_params(0x100000, true)
    }

    /// Parse a list of strings into chunks to be assembled.
    /// Each string may be a filename or a filename@offset describing where in the assembled image the contents of the file should appear.
    /// The offset is an integer expressed in any of the bases accepted by [`ParseInt`].
    pub fn parse(&mut self, chunks: &[impl AsRef<str>]) -> Result<()> {
        for chunk in chunks {
            if let Some((file, offset)) = chunk.as_ref().split_once('@') {
                self.chunks.push(ImageChunk::Offset(
                    PathBuf::from(file),
                    usize::from_str(offset)?,
                ));
            } else {
                self.chunks
                    .push(ImageChunk::Concat(PathBuf::from(chunk.as_ref())));
            }
        }
        Ok(())
    }

    // Read a file into a buffer.  Ensure the entire file is read.
    fn read(path: &Path, buf: &mut [u8]) -> Result<usize> {
        let mut file = File::open(path)?;
        let len = file.metadata()?.len() as usize;
        let n = file.read(buf)?;
        ensure!(len == n, ImageError::IncompleteRead(len, n));
        Ok(n)
    }

    /// Assemble the image according to the parameters and parsed chunk specifications.
    pub fn assemble(&self) -> Result<Vec<u8>> {
        let size = if self.mirrored {
            self.size / 2
        } else {
            self.size
        };
        let mut image = vec![0xff; size];
        let mut pos = 0;
        for chunk in &self.chunks {
            match chunk {
                ImageChunk::Concat(path) => {
                    let n = Self::read(path, &mut image[pos..])?;
                    pos += n;
                }
                ImageChunk::Offset(path, offset) => {
                    let n = Self::read(path, &mut image[*offset..])?;
                    pos = offset + n;
                }
            }
        }
        if self.mirrored {
            image.extend_from_within(..size);
        }
        Ok(image)
    }
}

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

    #[test]
    fn test_assemble_concat() -> Result<()> {
        // Test image assembly by concatenation.
        let mut image = ImageAssembler::with_params(16, false);
        image.parse(&[
            testdata!("hello.txt").to_str().unwrap(),
            testdata!("world.txt").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"HelloWorld\xff\xff\xff\xff\xff\xff");
        Ok(())
    }

    #[test]
    fn test_assemble_offset() -> Result<()> {
        // Test image assembly by explicit offsets.
        let mut image = ImageAssembler::with_params(16, false);
        image.parse(&[
            testdata!("hello.txt@0").to_str().unwrap(),
            testdata!("world.txt@0x8").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"Hello\xff\xff\xffWorld\xff\xff\xff");
        Ok(())
    }

    #[test]
    fn test_assemble_mirrored() -> Result<()> {
        // Test image assembly with mirroring.
        let mut image = ImageAssembler::with_params(20, true);
        image.parse(&[
            testdata!("hello.txt").to_str().unwrap(),
            testdata!("world.txt").to_str().unwrap(),
        ])?;
        let data = image.assemble()?;
        assert_eq!(data, b"HelloWorldHelloWorld");
        Ok(())
    }

    #[test]
    fn test_assemble_mirrored_offset_error() -> Result<()> {
        // Test image assembly where one of the source files isn't read completely.
        let mut image = ImageAssembler::with_params(16, true);
        image.parse(&[
            testdata!("hello.txt@0").to_str().unwrap(),
            testdata!("world.txt@0x5").to_str().unwrap(),
        ])?;
        let err = image.assemble().unwrap_err();
        assert_eq!(
            err.to_string(),
            "Incomplete read: expected to read 5 bytes but read 3 bytes"
        );
        Ok(())
    }

    #[test]
    fn test_load_image() {
        // Read and write back image.
        let image = Image::read_from_file(&testdata!("test_image.bin")).unwrap();
        image
            .write_to_file(&testdata!("test_image_out.bin"))
            .unwrap();

        // Ensure the result is identical to the original.
        let (mut orig_bytes, mut res_bytes) = (Vec::<u8>::new(), Vec::<u8>::new());
        File::open(&testdata!("test_image.bin"))
            .unwrap()
            .read_to_end(&mut orig_bytes)
            .unwrap();
        File::open(&testdata!("test_image_out.bin"))
            .unwrap()
            .read_to_end(&mut res_bytes)
            .unwrap();
        assert_eq!(orig_bytes, res_bytes);
    }
}
<|endoftext|>


<commit_before>/// Figure 10.22 Protecting a critical region from a signal
///
/// Takeaway: this is the first place where to_option() is wrong, as we
/// explicetly expect a -1 from sigsuspend

extern crate libc;
extern crate apue;
extern crate errno;

use libc::{c_int, SIGINT, SIG_ERR, SIGUSR1, SIG_BLOCK, SIG_SETMASK};
use libc::{signal, sigemptyset, sigaddset};
use apue::my_libc::{sigprocmask, sigsuspend};
use apue::{pr_mask, LibcResult};
use std::mem::uninitialized as uninit;

fn sig_int(_: c_int) {
    pr_mask("\nin sig_int: ");
}

fn main() {
    unsafe {
        let (mut waitmask, mut newmask, mut oldmask) = (uninit(), uninit(), uninit());
        pr_mask("program start: ");
        if signal(SIGINT, sig_int as usize) == SIG_ERR {
            println!("signal(SIGINT) error");
        }
        sigemptyset(&mut waitmask);
        sigaddset(&mut waitmask, SIGUSR1);
        sigemptyset(&mut newmask);
        sigaddset(&mut newmask, SIGINT);

        // Block SIGINT and save current signal mask
        sigprocmask(SIG_BLOCK, &newmask, &mut oldmask).to_option().expect("SIG_BLOCK error");
        // Critical region of code
        pr_mask("in critical region: ");
        // PAuse, allowing all signals except SIGUSR1
        if sigsuspend(&waitmask) != -1 {
            panic!("sigsuspend error");
        }
        pr_mask("after return from sigsuspend: ");
        // reset signal mask which unblocks SIGINT
        sigprocmask(SIG_SETMASK, &oldmask, std::ptr::null_mut())
            .to_option()
            .expect("SIG_SETMASK error");
        pr_mask("program exit: ");
    }
}
<commit_msg>fix typo<commit_after>/// Figure 10.22 Protecting a critical region from a signal
///
/// Takeaway: this is the first place where to_option() is wrong, as we
/// explicetly expect a -1 from sigsuspend

extern crate libc;
extern crate apue;
extern crate errno;

use libc::{c_int, SIGINT, SIG_ERR, SIGUSR1, SIG_BLOCK, SIG_SETMASK};
use libc::{signal, sigemptyset, sigaddset};
use apue::my_libc::{sigprocmask, sigsuspend};
use apue::{pr_mask, LibcResult};
use std::mem::uninitialized as uninit;

fn sig_int(_: c_int) {
    pr_mask("\nin sig_int: ");
}

fn main() {
    unsafe {
        let (mut waitmask, mut newmask, mut oldmask) = (uninit(), uninit(), uninit());
        pr_mask("program start: ");
        if signal(SIGINT, sig_int as usize) == SIG_ERR {
            println!("signal(SIGINT) error");
        }
        sigemptyset(&mut waitmask);
        sigaddset(&mut waitmask, SIGUSR1);
        sigemptyset(&mut newmask);
        sigaddset(&mut newmask, SIGINT);

        // Block SIGINT and save current signal mask
        sigprocmask(SIG_BLOCK, &newmask, &mut oldmask).to_option().expect("SIG_BLOCK error");
        // Critical region of code
        pr_mask("in critical region: ");
        // Pause, allowing all signals except SIGUSR1
        if sigsuspend(&waitmask) != -1 {
            panic!("sigsuspend error");
        }
        pr_mask("after return from sigsuspend: ");
        // reset signal mask which unblocks SIGINT
        sigprocmask(SIG_SETMASK, &oldmask, std::ptr::null_mut())
            .to_option()
            .expect("SIG_SETMASK error");
        pr_mask("program exit: ");
    }
}<|endoftext|>