<commit_before>//! The YAML parser. Here is where we convert the yaml in to commands to be
/// processed later.

use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Attach, KillWindow};
use rand::random;

#[cfg(test)] use yaml_rust::{YamlLoader};

/// Here was pass in the parsed yaml and project name. The purpose of this main
/// loop is to build the stack of commands that are run to setup a users tmux
/// session.
///
/// yaml_string: The parsed yaml from the config file.
/// project_name: The name of the project.
pub fn main(yaml_string: &Vec<Yaml>, project_name: &String, daemonize: bool) -> Result<Vec<Command>, String> {
    let mut commands: Vec<Command> = vec!();
    let tmp_window_name = format!("muxed_first_window_{}", random::<u16>());

    // The initial session command. Contains the tmp_window_name to be closed
    // before attaching.
    commands.push(Command::Session(Session{name: project_name.clone(), tmp_window_name: tmp_window_name.clone()}));

    // There should only be one doc but it's a vec so loop it.
    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let windows = doc["windows"].as_vec().expect("No Windows have been defined.");
        for (i, window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: k.as_str().unwrap().to_string(), root: root.clone()}));
                            commands.append(&mut try!(pane_matcher(&project_name, v, &root, k.as_str().unwrap().to_string())));
                        } else {
                            commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: try!(k.as_str().ok_or_else(|| "Windows require being named in your config.").map(|x| x.to_string())), root: root.clone()}));

                            if v.as_str().is_some() {
                                commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(), exec: v.as_str().unwrap().to_string()}));
                            }
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.clone(), root: root.clone()}))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{session_name: format!("{}:{}", project_name, i+1), name: s.to_string(), root: root.clone()}))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

    commands.push(Command::KillWindow(KillWindow{name: format!("{}:{}", project_name, tmp_window_name)}));

    if !daemonize { commands.push(Command::Attach(Attach{name: project_name.clone()})) };
    Ok(commands)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher(session: &String, window: &Yaml, root: &Option<String>, window_name: String) -> Result<Vec<Command>, String> {
    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        // For every pane, we need one less split.
        // ex. An existing window to become 2 panes, needs 1 split.
        if i < (panes.len()-1) {
            commands.push(Command::Split(Split{target: format!("{}:{}.{}", session, window_name, i).to_string(), root: root.clone()}));
        };
        // Execute given commands in each new pane after all splits are
        // complete.

        if pane.as_str().is_some() {
            commands.push(Command::SendKeys(SendKeys{target: format!("{}:{}.{}", session, window_name, i).to_string(), exec: pane.as_str().unwrap().to_string()}));
        }
    };

    // After all panes are split select the layout for the window
    if window["layout"].as_str().is_some() {
        let layout = window["layout"].as_str().expect("Bad layout").to_string();
        commands.push(Command::Layout(Layout{target: format!("{}:{}", session, window_name).to_string(), layout: layout}));
    };

    Ok(commands)
}

#[test]
pub fn windows_defined_as_array_has_6_commands() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_defined_as_array_has_1_session() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_defined_as_array_has_3_windows() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 3)
}

#[test]
pub fn windows_defined_as_array_has_1_attach() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Attach(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 9)
}

#[test]
pub fn windows_with_empty_command() {
    let s = "---
windows:
  - editor:
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn window_with_empty_command_has_no_send_keys() {
    let s = "---
windows:
  - editor:
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::SendKeys(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 0)
}

#[test]
pub fn window_with_no_name() {
    let s = "---
windows:
  - : ls
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_err())
}

#[test]
pub fn panes_with_empty_command() {
    let s = "---
windows:
  - cargo:
      layout: 'main-vertical'
      panes:
        -
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn panes_with_empty_command_has_no_send_keys() {
    let s = "---
windows:
  - editor:
      panes:
        -
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::SendKeys(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 0)
}

#[test]
pub fn panes_array_has_8_commands() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 8)
}

#[test]
pub fn panes_array_has_1_split() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Split(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_layout() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Layout(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_window() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_session() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}
<commit_msg>Make shorter lines.<commit_after>//! The YAML parser. Here is where we convert the yaml in to commands to be
/// processed later.

use yaml_rust::Yaml;
use command::{Command, Session, SendKeys, Split, Layout, Window, Attach, KillWindow};
use rand::random;

#[cfg(test)] use yaml_rust::{YamlLoader};

/// Here was pass in the parsed yaml and project name. The purpose of this main
/// loop is to build the stack of commands that are run to setup a users tmux
/// session.
///
/// yaml_string: The parsed yaml from the config file.
/// project_name: The name of the project.
pub fn main(yaml_string: &Vec<Yaml>, project_name: &String, daemonize: bool) -> Result<Vec<Command>, String> {
    let mut commands: Vec<Command> = vec!();
    let tmp_window_name = format!("muxed_first_window_{}", random::<u16>());

    // The initial session command. Contains the tmp_window_name to be closed
    // before attaching.
    commands.push(Command::Session(Session{
        name: project_name.clone(),
        tmp_window_name: tmp_window_name.clone()
    }));

    // There should only be one doc but it's a vec so loop it.
    for doc in yaml_string {
        let root = match doc["root"].as_str() {
            Some(x) => Some(x.to_string()),
            None    => None
        };

        let windows = doc["windows"].as_vec().expect("No Windows have been defined.");
        for (i, window) in windows.iter().enumerate() {
            match window {
                &Yaml::Hash(ref h)  => {
                    for (k, v) in h {
                        if v.as_hash().is_some() {
                            commands.push(Command::Window(Window{
                                    session_name: format!("{}:{}", project_name, i+1),
                                    name: k.as_str().unwrap().to_string(),
                                    root: root.clone()
                            }));

                            commands.append(&mut try!(pane_matcher(&project_name, v, &root, k.as_str().unwrap().to_string())));
                        } else {
                            commands.push(Command::Window(Window{
                                session_name: format!("{}:{}", project_name, i+1),
                                name: try!(k.as_str().ok_or_else(|| "Windows require being named in your config.").map(|x| x.to_string())),
                                root: root.clone()
                            }));

                            if v.as_str().is_some() {
                                commands.push(Command::SendKeys(SendKeys{
                                    target: format!("{}:{}", project_name, k.as_str().unwrap().to_string()).to_string(),
                                    exec: v.as_str().unwrap().to_string()
                                }));
                            }
                        }
                    }
                },
                &Yaml::String(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: format!("{}:{}", project_name, i+1),
                        name: s.clone(),
                        root: root.clone()
                    }))
                },
                &Yaml::Integer(ref s) => {
                    commands.push(Command::Window(Window{
                        session_name: format!("{}:{}", project_name, i+1),
                        name: s.to_string(),
                        root: root.clone()
                    }))
                },
                _ => panic!("Muxed config file formatting isn't recognized.")
            };
        };
    };

    commands.push(Command::KillWindow(KillWindow{name: format!("{}:{}", project_name, tmp_window_name)}));

    if !daemonize { commands.push(Command::Attach(Attach{name: project_name.clone()})) };
    Ok(commands)
}

/// Pane matcher is for breaking apart the panes. Splitting windows when needed
/// and executing commands as needed.
fn pane_matcher(session: &String, window: &Yaml, root: &Option<String>, window_name: String) -> Result<Vec<Command>, String> {
    let mut commands = vec!();
    let panes = window["panes"].as_vec().expect("Something is wrong with panes.");

    for (i, pane) in panes.iter().enumerate() {
        // For every pane, we need one less split.
        // ex. An existing window to become 2 panes, needs 1 split.
        if i < (panes.len()-1) {
            commands.push(Command::Split(Split{
                target: format!("{}:{}.{}", session, window_name, i).to_string(),
                root: root.clone()
            }));
        };
        // Execute given commands in each new pane after all splits are
        // complete.

        if pane.as_str().is_some() {
            commands.push(Command::SendKeys(SendKeys{
                target: format!("{}:{}.{}", session, window_name, i).to_string(),
                exec: pane.as_str().unwrap().to_string()
            }));
        }
    };

    // After all panes are split select the layout for the window
    if window["layout"].as_str().is_some() {
        let layout = window["layout"].as_str().expect("Bad layout").to_string();
        commands.push(Command::Layout(Layout{
            target: format!("{}:{}", session, window_name).to_string(),
            layout: layout
        }));
    };

    Ok(commands)
}

#[test]
pub fn windows_defined_as_array_has_6_commands() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_defined_as_array_has_1_session() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_defined_as_array_has_3_windows() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 3)
}

#[test]
pub fn windows_defined_as_array_has_1_attach() {
    let s = "---
windows: ['cargo', 'vim', 'git']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Attach(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn windows_with_integer_names() {
    let s = "---
windows: [1, 'vim', 3]
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    assert_eq!(main(&yaml, &"muxed".to_string(), false).unwrap().len(), 6)
}

#[test]
pub fn windows_as_list() {
    let s = "---
windows:
  - cargo: ''
  - vim: ''
  - git: ''
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 9)
}

#[test]
pub fn windows_with_empty_command() {
    let s = "---
windows:
  - editor:
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn window_with_empty_command_has_no_send_keys() {
    let s = "---
windows:
  - editor:
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::SendKeys(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 0)
}

#[test]
pub fn window_with_no_name() {
    let s = "---
windows:
  - : ls
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_err())
}

#[test]
pub fn panes_with_empty_command() {
    let s = "---
windows:
  - cargo:
      layout: 'main-vertical'
      panes:
        -
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let result = main(&yaml, &"muxed".to_string(), false);
    assert!(result.is_ok())
}

#[test]
pub fn panes_with_empty_command_has_no_send_keys() {
    let s = "---
windows:
  - editor:
      panes:
        -
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::SendKeys(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 0)
}

#[test]
pub fn panes_array_has_8_commands() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";
    let yaml = YamlLoader::load_from_str(s).unwrap();
    let commands = main(&yaml, &"muxed".to_string(), false).unwrap();
    assert_eq!(commands.len(), 8)
}

#[test]
pub fn panes_array_has_1_split() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Split(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_layout() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Layout(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_window() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Window(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}

#[test]
pub fn panes_array_has_1_session() {
    let s = "---
windows:
  - editor:
      layout: 'main-vertical'
      panes: ['vim', 'guard']
";

    let yaml = YamlLoader::load_from_str(s).unwrap();
    let remains: Vec<Command> = main(&yaml, &"muxed".to_string(), false).unwrap().into_iter().filter(|x| match x {
        &Command::Session(_) => true,
        _ => false
    }).collect();

    assert_eq!(remains.len(), 1)
}
<|endoftext|>


<commit_before>use satyrs::cnf::{CNF, Assignment, PartialAssignment};
use std::io;

/**
 * DPLL Algorithm implementation
 * Start by cloning the cnf so we can modify the HashMap
 * Clone might be slow, consider a faster option
 * TODO: Consider using HashSet instead of Vec for the clauses
 */
#[allow(non_snake_case)]
pub fn DPLL(cnf: &CNF) -> Option<Assignment> {
    //let _cnf = cnf.clone();
    //println!("DPLL {:?}", _cnf);
    let mut p_assn = PartialAssignment::new(cnf.nvar as usize);
    match _dpll(cnf, &mut p_assn) {
        Some(assn) => Some(assn.assignment.iter().map(|a| {
            println!("WE HERE ");
            match *a {
                Some(a) => a,
                None => true,
            }
        }).collect()),
        None => None,
    }
}

#[allow(unused_variables)]
fn _dpll(cnf: &CNF, p_assn: &mut PartialAssignment) -> Option<PartialAssignment> {
    // If consistent set of literals, return True
    println!("{}\n{}",cnf,p_assn);
    if cnf.clauses.is_empty() {
        let _p_assn = p_assn.clone();
        return Some(_p_assn);  // Display optional value
    }

    // If contains an empty clause return None
    // TODO: Consider optimizing
    for clause in cnf.clauses.values() {
        if clause.is_empty() {
            return None;
        }
    }

    // For every unit-clause, unit-propogate
    let mut _cnf = cnf.clone();
    for unit in cnf.units.iter() {
        // XXX: Previous versions of this for loop could have unit-propagated and
        // remove the unit clause here, so it's not necessary that the unit clause id
        // still exists in the clauses.
        if let Some(clause) = cnf.clauses.get(&unit) {
            // Then via unit propagation we've created an empty clause; no solution down this path
            if clause.is_empty() {
                return None;
            }
            let lit : i32 = zeroth!(clause);
            p_assn.assign_literal(lit);
            _cnf.unit_propagate(*unit); // Only propogate in the clone
        }
    } // FIXME: Unit propogate might be the end of things
    println!("After Prop! {}",_cnf);
    // Clone for the right
    let mut r_cnf = _cnf.clone();

    // TODO: For every pure literal, pure literal assign
    // Choose literal L for split
    //let literal: i32 = zeroth!(_cnf.clauses.values().next().unwrap());
    let literal = match _cnf.clauses.values().next() {
        Some(clause) => {
            if !clause.is_empty() {
                Some(zeroth!(clause))
            } else {
                return None;
            }
        }
        None => None,
    };

    if literal.is_none() {
        let _p_assn = p_assn.clone();
        return Some(_p_assn);
    }
    let lit = literal.unwrap();
    println!("Split on {} ({})",lit,lit/2);
    // Propagate literal
    _cnf.propagate(lit);

    // Return DPLL with L and -L
    p_assn.assign_literal(lit);

    //let mut buffer = String::new();
    //let _ = io::stdin().read_line(&mut buffer);

    let left = _dpll(&_cnf, p_assn);
    if left.is_some() { return left; }
    p_assn.unassign_literal(lit);
    p_assn.assign_literal(lit ^ 1);
    r_cnf.propagate(lit^1);
    _dpll(&_cnf, p_assn)
}
<commit_msg>Fix dpll bug; pass in correct version of formula<commit_after>use satyrs::cnf::{CNF, Assignment, PartialAssignment};
use std::io;

/**
 * DPLL Algorithm implementation
 * Start by cloning the cnf so we can modify the HashMap
 * Clone might be slow, consider a faster option
 * TODO: Consider using HashSet instead of Vec for the clauses
 */
#[allow(non_snake_case)]
pub fn DPLL(cnf: &CNF) -> Option<Assignment> {
    //let _cnf = cnf.clone();
    //println!("DPLL {:?}", _cnf);
    let mut p_assn = PartialAssignment::new(cnf.nvar as usize);
    match _dpll(cnf, &mut p_assn) {
        Some(assn) => Some(assn.assignment.iter().map(|a| {
            println!("WE HERE ");
            match *a {
                Some(a) => a,
                None => true,
            }
        }).collect()),
        None => None,
    }
}

#[allow(unused_variables)]
fn _dpll(cnf: &CNF, p_assn: &mut PartialAssignment) -> Option<PartialAssignment> {
    // If consistent set of literals, return True
    println!("{}\n{}",cnf,p_assn);
    if cnf.clauses.is_empty() {
        let _p_assn = p_assn.clone();
        return Some(_p_assn);  // Display optional value
    }

    // If contains an empty clause return None
    // TODO: Consider optimizing
    for clause in cnf.clauses.values() {
        if clause.is_empty() {
            return None;
        }
    }

    // For every unit-clause, unit-propogate
    let mut _cnf = cnf.clone();
    for unit in cnf.units.iter() {
        // XXX: Previous versions of this for loop could have unit-propagated and
        // remove the unit clause here, so it's not necessary that the unit clause id
        // still exists in the clauses.
        if let Some(clause) = cnf.clauses.get(&unit) {
            // Then via unit propagation we've created an empty clause; no solution down this path
            if clause.is_empty() {
                return None;
            }
            let lit : i32 = zeroth!(clause);
            p_assn.assign_literal(lit);
            _cnf.unit_propagate(*unit); // Only propogate in the clone
        }
    } // FIXME: Unit propogate might be the end of things
    println!("After Prop! {}",_cnf);
    // Clone for the right
    let mut r_cnf = _cnf.clone();

    // TODO: For every pure literal, pure literal assign
    // Choose literal L for split
    //let literal: i32 = zeroth!(_cnf.clauses.values().next().unwrap());
    let literal = match _cnf.clauses.values().next() {
        Some(clause) => {
            if !clause.is_empty() {
                Some(zeroth!(clause))
            } else {
                return None;
            }
        }
        None => None,
    };

    if literal.is_none() {
        let _p_assn = p_assn.clone();
        return Some(_p_assn);
    }
    let lit = literal.unwrap();
    println!("Split on {} ({})",lit,lit/2);
    // Propagate literal
    _cnf.propagate(lit);

    // Return DPLL with L and -L
    p_assn.assign_literal(lit);

    //let mut buffer = String::new();
    //let _ = io::stdin().read_line(&mut buffer);

    let left = _dpll(&_cnf, p_assn);
    if left.is_some() { return left; }
    p_assn.unassign_literal(lit);
    p_assn.assign_literal(lit ^ 1);
    r_cnf.propagate(lit^1);
    _dpll(&r_cnf, p_assn)
}
<|endoftext|>


<commit_before>use std::path::PathBuf;

pub enum GameState {
    NotStarted,
    Running,
    Paused,
}

pub struct Context {
    pub assets_path: PathBuf,
    pub game_state: GameState,
}
<commit_msg>Paused state has been removed<commit_after>use std::path::PathBuf;

pub enum GameState {
    NotStarted,
    Started,
}

pub struct Context {
    pub assets_path: PathBuf,
    pub game_state: GameState,
}
<|endoftext|>


<commit_before>use gfx;

gfx_defines! {
  constant TileMapData {
    data: [f32; 4] = "data",
  }

  constant TilemapSettings {
    world_size: [f32; 2] = "u_WorldSize",
    tilesheet_size: [f32; 2] = "u_TilesheetSize",
  }

  vertex VertexData {
    pos: [f32; 2] = "a_Pos",
    buf_pos: [f32; 2] = "a_BufPos",
  }

  constant CharacterSheet {
    x_div: f32 = "x_div",
    y_div: f32 = "y_div",
    row_idx: u32 = "a_row",
    index: f32 = "a_index",
  }

  constant Position {
    position: [f32; 2] = "a_position",
  }

  pipeline bullet_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    position_cb: gfx::ConstantBuffer<Position> = "b_BulletPosition",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline critter_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    position_cb: gfx::ConstantBuffer<Position> = "b_CharacterPosition",
    character_sprite_cb: gfx::ConstantBuffer<CharacterSheet> = "b_CharacterSprite",
    charactersheet: gfx::TextureSampler<[f32; 4]> = "t_CharacterSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline tilemap_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_TileMapPosition",
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    tilemap: gfx::ConstantBuffer<TileMapData> = "b_TileMap",
    tilemap_cb: gfx::ConstantBuffer<TilemapSettings> = "b_PsLocals",
    tilesheet: gfx::TextureSampler<[f32; 4]> = "t_TileSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline static_element_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_StaticElementPosition",
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    static_element_sheet: gfx::TextureSampler<[f32; 4]> = "t_StaticElementSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline text_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_TextPosition",
    text_sheet: gfx::TextureSampler<[f32; 4]> = "t_TextSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  constant Projection {
    model: [[f32; 4]; 4] = "u_Model",
    view: [[f32; 4]; 4] = "u_View",
    proj: [[f32; 4]; 4] = "u_Proj",
  }
}

impl VertexData {
  pub fn new(pos: [f32; 2], buf_pos: [f32; 2]) -> VertexData {
    VertexData {
      pos,
      buf_pos,
    }
  }
}

impl TileMapData {
  pub fn new_empty() -> TileMapData {
    TileMapData { data: [32.0, 32.0, 0.0, 0.0] }
  }

  pub fn new(data: [f32; 4]) -> TileMapData {
    TileMapData { data }
  }
}

impl Position {
  pub fn new(position: [f32; 2]) -> Position { Position { position } }
}
<commit_msg>Implement Add for Position type<commit_after>use gfx;
use std::ops::Add;

gfx_defines! {
  constant TileMapData {
    data: [f32; 4] = "data",
  }

  constant TilemapSettings {
    world_size: [f32; 2] = "u_WorldSize",
    tilesheet_size: [f32; 2] = "u_TilesheetSize",
  }

  vertex VertexData {
    pos: [f32; 2] = "a_Pos",
    buf_pos: [f32; 2] = "a_BufPos",
  }

  constant CharacterSheet {
    x_div: f32 = "x_div",
    y_div: f32 = "y_div",
    row_idx: u32 = "a_row",
    index: f32 = "a_index",
  }

  constant Position {
    position: [f32; 2] = "a_position",
  }

  pipeline bullet_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    position_cb: gfx::ConstantBuffer<Position> = "b_BulletPosition",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline critter_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    position_cb: gfx::ConstantBuffer<Position> = "b_CharacterPosition",
    character_sprite_cb: gfx::ConstantBuffer<CharacterSheet> = "b_CharacterSprite",
    charactersheet: gfx::TextureSampler<[f32; 4]> = "t_CharacterSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline tilemap_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_TileMapPosition",
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    tilemap: gfx::ConstantBuffer<TileMapData> = "b_TileMap",
    tilemap_cb: gfx::ConstantBuffer<TilemapSettings> = "b_PsLocals",
    tilesheet: gfx::TextureSampler<[f32; 4]> = "t_TileSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline static_element_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_StaticElementPosition",
    projection_cb: gfx::ConstantBuffer<Projection> = "b_VsLocals",
    static_element_sheet: gfx::TextureSampler<[f32; 4]> = "t_StaticElementSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  pipeline text_pipeline {
    vbuf: gfx::VertexBuffer<VertexData> = (),
    position_cb: gfx::ConstantBuffer<Position> = "b_TextPosition",
    text_sheet: gfx::TextureSampler<[f32; 4]> = "t_TextSheet",
    out_color: gfx::RenderTarget<gfx::format::Rgba8> = "Target0",
    out_depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
  }

  constant Projection {
    model: [[f32; 4]; 4] = "u_Model",
    view: [[f32; 4]; 4] = "u_View",
    proj: [[f32; 4]; 4] = "u_Proj",
  }
}

impl VertexData {
  pub fn new(pos: [f32; 2], buf_pos: [f32; 2]) -> VertexData {
    VertexData {
      pos,
      buf_pos,
    }
  }
}

impl TileMapData {
  pub fn new_empty() -> TileMapData {
    TileMapData { data: [32.0, 32.0, 0.0, 0.0] }
  }

  pub fn new(data: [f32; 4]) -> TileMapData {
    TileMapData { data }
  }
}

impl Position {
  pub fn new(position: [f32; 2]) -> Position { Position { position } }
}

impl Add for Position {
  type Output = Position;

  fn add(self, other: Position) -> Position {
    Position::new([self.position[0] + other.position[0], self.position[1] + other.position[1]])
  }
}
<|endoftext|>


<commit_before>//! This module provides the interface to different solvers.
//!
//! Both [`coin_cbc`](https://docs.rs/coin_cbc/latest/coin_cbc/) and
//! [`minilp`](https://docs.rs/minilp/0.2.2/minilp/) are available as cargo
//! [features](https://doc.rust-lang.org/cargo/reference/features.html). To use
//! them, specify your dependency to `lp_modeler` accordingly in your `Cargo.toml`
//! (note the name difference of the `native_coin_cbc` feature for the `coin_cbc` crate):
//! ```ignore
//! [dependencies.lp_modeler]
//! version = "4.3"
//! features = "native_coin_cbc"
//! ```
//! or:
//! ```ignore
//! [dependencies.lp_modeler]
//! version = "4.3"
//! features = "minilp"
//! ```
//! For `coin_cbc` to compile, the `Cbc` library files need to be available on your system.
//! See the [`coin_cbc` project README](https://github.com/KardinalAI/coin_cbc) for more infos.
//!
//! The other solvers need to be installed externally on your system.
//! The respective information is provided in the project's README in the section on
//! [installing external solvers](https://github.com/jcavat/rust-lp-modeler#installing-external-solvers).

use std::collections::HashMap;

use dsl::{Problem, LpContinuous, LpBinary, LpInteger, LpProblem, LpExprNode, LpExprOp, LpExprArenaIndex};

pub mod cbc;
pub use self::cbc::*;

pub mod gurobi;
pub use self::gurobi::*;

pub mod glpk;
pub use self::glpk::*;

#[cfg(feature = "minilp")]
pub mod minilp;
#[cfg(feature = "minilp")]
pub use self::minilp::*;

#[cfg(feature = "native_coin_cbc")]
pub mod native_cbc;
#[cfg(feature = "native_coin_cbc")]
pub use self::native_cbc::*;

use std::fs::File;
use std::fs;
use util::is_zero;

#[derive(Debug, PartialEq, Clone)]
pub enum Status {
    Optimal,
    SubOptimal,
    Infeasible,
    Unbounded,
    NotSolved,
}

#[derive(Debug, Clone)]
pub struct Solution<'a> {
    pub status: Status,
    pub results: HashMap<String, f32>,
    pub related_problem: Option<&'a LpProblem>
}
impl Solution<'_> {
    pub fn new<'a>(status: Status, results: HashMap<String, f32>) -> Solution<'a> {
        Solution {
            status,
            results,
            related_problem: None
        }
    }
    pub fn with_problem(status: Status, results: HashMap<String, f32>, problem: &LpProblem) -> Solution {
        Solution {
            status,
            results,
            related_problem: Some(problem)
        }
    }
    fn check_possible_solution(&self) {
        match &self.status {
            Status::Unbounded | Status::NotSolved | Status::Infeasible => panic!("Solution must be optimal or suboptimal"),
            _ => ()
        }
    }
    pub fn get_raw_value(&self, name: &str) -> f32 {
        self.check_possible_solution();
        *self.results.get(name).expect("No value found for this variable. Check if the variable has been used in the related problem.")
    }
    pub fn get_bool(&self, var: &LpBinary) -> bool {
        self.check_possible_solution();
        self.results.get(&var.name).and_then(|&f| if is_zero(1.0-f) { Some(true) } else if is_zero(f) { Some(false) } else { None } ).expect("Result value cannot be interpreted as boolean")
    }
    pub fn get_float(&self, var: &LpContinuous) -> f32 {
        self.check_possible_solution();
        *self.results.get(&var.name).expect("No value found for this variable. Check if the variable has been used in the related problem.")
    }
    pub fn get_int(&self, var: &LpInteger) -> i32 {
        self.check_possible_solution();
        let &f = self.results.get(&var.name).expect("No value found for this variable. Check if the variable has been used in the related problem.");
        let i = f as i32;
        assert!( is_zero( f-(i as f32)), format!("Value {} cannot be interpreted as integer.", f) );
        i
    }
    pub fn eval(&self) -> Option<f32> {
        self.related_problem
            .and_then( |problem| {
                match &problem.obj_expr_arena {
                    Some(obj_expr_arena) => Some( self.eval_with(&obj_expr_arena.get_root_index(), &self.results) ),
                    None => None
                }
            })
    }
    fn eval_with(&self, index: &LpExprArenaIndex, values: &HashMap<String, f32>) -> f32 {
        match self.related_problem.unwrap().obj_expr_arena.as_ref().unwrap().expr_ref_at(*index) {
            LpExprNode::LpCompExpr(operation, left, right) => {
                match operation {
                    LpExprOp::Addition => self.eval_with(left, values) + self.eval_with(right, values),
                    LpExprOp::Multiplication => self.eval_with(left, values) * self.eval_with(right, values),
                    LpExprOp::Subtraction => self.eval_with(left, values) - self.eval_with(right, values),
                }
            },
            LpExprNode::ConsBin(LpBinary { name })
            | LpExprNode::ConsCont(LpContinuous { name, .. })
            | LpExprNode::ConsInt(LpInteger { name, .. }) => *values.get(name).unwrap_or(&0f32),
            LpExprNode::LitVal(n) => *n,
            LpExprNode::EmptyExpr => 0.0
        }
    }
}

pub trait SolverTrait {
    type P: Problem;
    fn run<'a>(&self, problem: &'a Self::P) -> Result<Solution<'a>, String>;
}

pub trait SolverWithSolutionParsing {
    fn read_solution<'a>(&self, temp_solution_file: &String, problem: Option<&'a LpProblem>) -> Result<Solution<'a>, String> {
        match File::open( temp_solution_file ) {
            Ok(f) => {
                let res = self.read_specific_solution(&f, problem)?;
                let _ = fs::remove_file(temp_solution_file);
                Ok(res)
            }
            Err(_) => return Err("Cannot open file".to_string()),
        }
    }
    fn read_specific_solution<'a>(&self, f: &File, problem: Option<&'a LpProblem>) -> Result<Solution<'a>, String>;
}

pub trait WithMaxSeconds<T> {
    fn max_seconds(&self) -> Option<u32>;
    fn with_max_seconds(&self, seconds: u32) -> T;
}

pub trait WithNbThreads<T> {
    fn nb_threads(&self) -> Option<u32>;
    fn with_nb_threads(&self, threads: u32) -> T;
}
<commit_msg>mark doctests with dep spec as ```toml instead of ```ignore<commit_after>//! This module provides the interface to different solvers.
//!
//! Both [`coin_cbc`](https://docs.rs/coin_cbc/latest/coin_cbc/) and
//! [`minilp`](https://docs.rs/minilp/0.2.2/minilp/) are available as cargo
//! [features](https://doc.rust-lang.org/cargo/reference/features.html). To use
//! them, specify your dependency to `lp_modeler` accordingly in your `Cargo.toml`
//! (note the name difference of the `native_coin_cbc` feature for the `coin_cbc` crate):
//! ```toml
//! [dependencies.lp_modeler]
//! version = "4.3"
//! features = "native_coin_cbc"
//! ```
//! or:
//! ```toml
//! [dependencies.lp_modeler]
//! version = "4.3"
//! features = "minilp"
//! ```
//! For `coin_cbc` to compile, the `Cbc` library files need to be available on your system.
//! See the [`coin_cbc` project README](https://github.com/KardinalAI/coin_cbc) for more infos.
//!
//! The other solvers need to be installed externally on your system.
//! The respective information is provided in the project's README in the section on
//! [installing external solvers](https://github.com/jcavat/rust-lp-modeler#installing-external-solvers).

use std::collections::HashMap;

use dsl::{Problem, LpContinuous, LpBinary, LpInteger, LpProblem, LpExprNode, LpExprOp, LpExprArenaIndex};

pub mod cbc;
pub use self::cbc::*;

pub mod gurobi;
pub use self::gurobi::*;

pub mod glpk;
pub use self::glpk::*;

#[cfg(feature = "minilp")]
pub mod minilp;
#[cfg(feature = "minilp")]
pub use self::minilp::*;

#[cfg(feature = "native_coin_cbc")]
pub mod native_cbc;
#[cfg(feature = "native_coin_cbc")]
pub use self::native_cbc::*;

use std::fs::File;
use std::fs;
use util::is_zero;

#[derive(Debug, PartialEq, Clone)]
pub enum Status {
    Optimal,
    SubOptimal,
    Infeasible,
    Unbounded,
    NotSolved,
}

#[derive(Debug, Clone)]
pub struct Solution<'a> {
    pub status: Status,
    pub results: HashMap<String, f32>,
    pub related_problem: Option<&'a LpProblem>
}
impl Solution<'_> {
    pub fn new<'a>(status: Status, results: HashMap<String, f32>) -> Solution<'a> {
        Solution {
            status,
            results,
            related_problem: None
        }
    }
    pub fn with_problem(status: Status, results: HashMap<String, f32>, problem: &LpProblem) -> Solution {
        Solution {
            status,
            results,
            related_problem: Some(problem)
        }
    }
    fn check_possible_solution(&self) {
        match &self.status {
            Status::Unbounded | Status::NotSolved | Status::Infeasible => panic!("Solution must be optimal or suboptimal"),
            _ => ()
        }
    }
    pub fn get_raw_value(&self, name: &str) -> f32 {
        self.check_possible_solution();
        *self.results.get(name).expect("No value found for this variable. Check if the variable has been used in the related problem.")
    }
    pub fn get_bool(&self, var: &LpBinary) -> bool {
        self.check_possible_solution();
        self.results.get(&var.name).and_then(|&f| if is_zero(1.0-f) { Some(true) } else if is_zero(f) { Some(false) } else { None } ).expect("Result value cannot be interpreted as boolean")
    }
    pub fn get_float(&self, var: &LpContinuous) -> f32 {
        self.check_possible_solution();
        *self.results.get(&var.name).expect("No value found for this variable. Check if the variable has been used in the related problem.")
    }
    pub fn get_int(&self, var: &LpInteger) -> i32 {
        self.check_possible_solution();
        let &f = self.results.get(&var.name).expect("No value found for this variable. Check if the variable has been used in the related problem.");
        let i = f as i32;
        assert!( is_zero( f-(i as f32)), format!("Value {} cannot be interpreted as integer.", f) );
        i
    }
    pub fn eval(&self) -> Option<f32> {
        self.related_problem
            .and_then( |problem| {
                match &problem.obj_expr_arena {
                    Some(obj_expr_arena) => Some( self.eval_with(&obj_expr_arena.get_root_index(), &self.results) ),
                    None => None
                }
            })
    }
    fn eval_with(&self, index: &LpExprArenaIndex, values: &HashMap<String, f32>) -> f32 {
        match self.related_problem.unwrap().obj_expr_arena.as_ref().unwrap().expr_ref_at(*index) {
            LpExprNode::LpCompExpr(operation, left, right) => {
                match operation {
                    LpExprOp::Addition => self.eval_with(left, values) + self.eval_with(right, values),
                    LpExprOp::Multiplication => self.eval_with(left, values) * self.eval_with(right, values),
                    LpExprOp::Subtraction => self.eval_with(left, values) - self.eval_with(right, values),
                }
            },
            LpExprNode::ConsBin(LpBinary { name })
            | LpExprNode::ConsCont(LpContinuous { name, .. })
            | LpExprNode::ConsInt(LpInteger { name, .. }) => *values.get(name).unwrap_or(&0f32),
            LpExprNode::LitVal(n) => *n,
            LpExprNode::EmptyExpr => 0.0
        }
    }
}

pub trait SolverTrait {
    type P: Problem;
    fn run<'a>(&self, problem: &'a Self::P) -> Result<Solution<'a>, String>;
}

pub trait SolverWithSolutionParsing {
    fn read_solution<'a>(&self, temp_solution_file: &String, problem: Option<&'a LpProblem>) -> Result<Solution<'a>, String> {
        match File::open( temp_solution_file ) {
            Ok(f) => {
                let res = self.read_specific_solution(&f, problem)?;
                let _ = fs::remove_file(temp_solution_file);
                Ok(res)
            }
            Err(_) => return Err("Cannot open file".to_string()),
        }
    }
    fn read_specific_solution<'a>(&self, f: &File, problem: Option<&'a LpProblem>) -> Result<Solution<'a>, String>;
}

pub trait WithMaxSeconds<T> {
    fn max_seconds(&self) -> Option<u32>;
    fn with_max_seconds(&self, seconds: u32) -> T;
}

pub trait WithNbThreads<T> {
    fn nb_threads(&self) -> Option<u32>;
    fn with_nb_threads(&self, threads: u32) -> T;
}
<|endoftext|>


<commit_before>//! Clippy wrappers around rustc's diagnostic functions.

use rustc_errors::{Applicability, DiagnosticBuilder};
use rustc_hir::HirId;
use rustc_lint::{LateContext, Lint, LintContext};
use rustc_span::source_map::{MultiSpan, Span};
use std::env;

fn docs_link(diag: &mut DiagnosticBuilder<'_>, lint: &'static Lint) {
    if env::var("CLIPPY_DISABLE_DOCS_LINKS").is_err() {
        diag.help(&format!(
            "for further information visit https://rust-lang.github.io/rust-clippy/{}/index.html#{}",
            &option_env!("RUST_RELEASE_NUM").map_or("master".to_string(), |n| {
                // extract just major + minor version and ignore patch versions
                format!("rust-{}", n.rsplitn(2, '.').nth(1).unwrap())
            }),
            lint.name_lower().replacen("clippy::", "", 1)
        ));
    }
}

/// Emit a basic lint message with a `msg` and a `span`.
///
/// This is the most primitive of our lint emission methods and can
/// be a good way to get a new lint started.
///
/// Usually it's nicer to provide more context for lint messages.
/// Be sure the output is understandable when you use this method.
///
/// # Example
///
/// ```ignore
/// error: usage of mem::forget on Drop type
///   --> $DIR/mem_forget.rs:17:5
///    |
/// 17 |     std::mem::forget(seven);
///    |     ^^^^^^^^^^^^^^^^^^^^^^^
/// ```
pub fn span_lint<T: LintContext>(cx: &T, lint: &'static Lint, sp: impl Into<MultiSpan>, msg: &str) {
    cx.struct_span_lint(lint, sp, |diag| {
        let mut diag = diag.build(msg);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Same as `span_lint` but with an extra `help` message.
///
/// Use this if you want to provide some general help but
/// can't provide a specific machine applicable suggestion.
///
/// The `help` message can be optionally attached to a `Span`.
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: constant division of 0.0 with 0.0 will always result in NaN
///   --> $DIR/zero_div_zero.rs:6:25
///    |
/// 6  |     let other_f64_nan = 0.0f64 / 0.0;
///    |                         ^^^^^^^^^^^^
///    |
///    = help: Consider using `f64::NAN` if you would like a constant representing NaN
/// ```
pub fn span_lint_and_help<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    span: Span,
    msg: &str,
    help_span: Option<Span>,
    help: &str,
) {
    cx.struct_span_lint(lint, span, |diag| {
        let mut diag = diag.build(msg);
        if let Some(help_span) = help_span {
            diag.span_help(help_span, help);
        } else {
            diag.help(help);
        }
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Like `span_lint` but with a `note` section instead of a `help` message.
///
/// The `note` message is presented separately from the main lint message
/// and is attached to a specific span:
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: calls to `std::mem::forget` with a reference instead of an owned value. Forgetting a reference does nothing.
///   --> $DIR/drop_forget_ref.rs:10:5
///    |
/// 10 |     forget(&SomeStruct);
///    |     ^^^^^^^^^^^^^^^^^^^
///    |
///    = note: `-D clippy::forget-ref` implied by `-D warnings`
/// note: argument has type &SomeStruct
///   --> $DIR/drop_forget_ref.rs:10:12
///    |
/// 10 |     forget(&SomeStruct);
///    |            ^^^^^^^^^^^
/// ```
pub fn span_lint_and_note<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    span: Span,
    msg: &str,
    note_span: Option<Span>,
    note: &str,
) {
    cx.struct_span_lint(lint, span, |diag| {
        let mut diag = diag.build(msg);
        if let Some(note_span) = note_span {
            diag.span_note(note_span, note);
        } else {
            diag.note(note);
        }
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Like `span_lint` but allows to add notes, help and suggestions using a closure.
///
/// If you need to customize your lint output a lot, use this function.
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
pub fn span_lint_and_then<'a, T: LintContext, F>(cx: &'a T, lint: &'static Lint, sp: Span, msg: &str, f: F)
where
    F: for<'b> FnOnce(&mut DiagnosticBuilder<'b>),
{
    cx.struct_span_lint(lint, sp, |diag| {
        let mut diag = diag.build(msg);
        f(&mut diag);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

pub fn span_lint_hir(cx: &LateContext<'_>, lint: &'static Lint, hir_id: HirId, sp: Span, msg: &str) {
    cx.tcx.struct_span_lint_hir(lint, hir_id, sp, |diag| {
        let mut diag = diag.build(msg);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

pub fn span_lint_hir_and_then(
    cx: &LateContext<'_>,
    lint: &'static Lint,
    hir_id: HirId,
    sp: Span,
    msg: &str,
    f: impl FnOnce(&mut DiagnosticBuilder<'_>),
) {
    cx.tcx.struct_span_lint_hir(lint, hir_id, sp, |diag| {
        let mut diag = diag.build(msg);
        f(&mut diag);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Add a span lint with a suggestion on how to fix it.
///
/// These suggestions can be parsed by rustfix to allow it to automatically fix your code.
/// In the example below, `help` is `"try"` and `sugg` is the suggested replacement `".any(|x| x >
/// 2)"`.
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: This `.fold` can be more succinctly expressed as `.any`
/// --> $DIR/methods.rs:390:13
///     |
/// 390 |     let _ = (0..3).fold(false, |acc, x| acc || x > 2);
///     |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: try: `.any(|x| x > 2)`
///     |
///     = note: `-D fold-any` implied by `-D warnings`
/// ```

#[allow(clippy::unknown_clippy_lints)]
#[cfg_attr(feature = "internal-lints", allow(clippy::collapsible_span_lint_calls))]
pub fn span_lint_and_sugg<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    sp: Span,
    msg: &str,
    help: &str,
    sugg: String,
    applicability: Applicability,
) {
    span_lint_and_then(cx, lint, sp, msg, |diag| {
        diag.span_suggestion(sp, help, sugg, applicability);
    });
}

/// Create a suggestion made from several `span → replacement`.
///
/// Note: in the JSON format (used by `compiletest_rs`), the help message will
/// appear once per
/// replacement. In human-readable format though, it only appears once before
/// the whole suggestion.
pub fn multispan_sugg<I>(diag: &mut DiagnosticBuilder<'_>, help_msg: &str, sugg: I)
where
    I: IntoIterator<Item = (Span, String)>,
{
    multispan_sugg_with_applicability(diag, help_msg, Applicability::Unspecified, sugg)
}

pub fn multispan_sugg_with_applicability<I>(
    diag: &mut DiagnosticBuilder<'_>,
    help_msg: &str,
    applicability: Applicability,
    sugg: I,
) where
    I: IntoIterator<Item = (Span, String)>,
{
    diag.multipart_suggestion(help_msg, sugg.into_iter().collect(), applicability);
}
<commit_msg>Remove unknown_clippy_lints allow attribute<commit_after>//! Clippy wrappers around rustc's diagnostic functions.

use rustc_errors::{Applicability, DiagnosticBuilder};
use rustc_hir::HirId;
use rustc_lint::{LateContext, Lint, LintContext};
use rustc_span::source_map::{MultiSpan, Span};
use std::env;

fn docs_link(diag: &mut DiagnosticBuilder<'_>, lint: &'static Lint) {
    if env::var("CLIPPY_DISABLE_DOCS_LINKS").is_err() {
        diag.help(&format!(
            "for further information visit https://rust-lang.github.io/rust-clippy/{}/index.html#{}",
            &option_env!("RUST_RELEASE_NUM").map_or("master".to_string(), |n| {
                // extract just major + minor version and ignore patch versions
                format!("rust-{}", n.rsplitn(2, '.').nth(1).unwrap())
            }),
            lint.name_lower().replacen("clippy::", "", 1)
        ));
    }
}

/// Emit a basic lint message with a `msg` and a `span`.
///
/// This is the most primitive of our lint emission methods and can
/// be a good way to get a new lint started.
///
/// Usually it's nicer to provide more context for lint messages.
/// Be sure the output is understandable when you use this method.
///
/// # Example
///
/// ```ignore
/// error: usage of mem::forget on Drop type
///   --> $DIR/mem_forget.rs:17:5
///    |
/// 17 |     std::mem::forget(seven);
///    |     ^^^^^^^^^^^^^^^^^^^^^^^
/// ```
pub fn span_lint<T: LintContext>(cx: &T, lint: &'static Lint, sp: impl Into<MultiSpan>, msg: &str) {
    cx.struct_span_lint(lint, sp, |diag| {
        let mut diag = diag.build(msg);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Same as `span_lint` but with an extra `help` message.
///
/// Use this if you want to provide some general help but
/// can't provide a specific machine applicable suggestion.
///
/// The `help` message can be optionally attached to a `Span`.
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: constant division of 0.0 with 0.0 will always result in NaN
///   --> $DIR/zero_div_zero.rs:6:25
///    |
/// 6  |     let other_f64_nan = 0.0f64 / 0.0;
///    |                         ^^^^^^^^^^^^
///    |
///    = help: Consider using `f64::NAN` if you would like a constant representing NaN
/// ```
pub fn span_lint_and_help<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    span: Span,
    msg: &str,
    help_span: Option<Span>,
    help: &str,
) {
    cx.struct_span_lint(lint, span, |diag| {
        let mut diag = diag.build(msg);
        if let Some(help_span) = help_span {
            diag.span_help(help_span, help);
        } else {
            diag.help(help);
        }
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Like `span_lint` but with a `note` section instead of a `help` message.
///
/// The `note` message is presented separately from the main lint message
/// and is attached to a specific span:
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: calls to `std::mem::forget` with a reference instead of an owned value. Forgetting a reference does nothing.
///   --> $DIR/drop_forget_ref.rs:10:5
///    |
/// 10 |     forget(&SomeStruct);
///    |     ^^^^^^^^^^^^^^^^^^^
///    |
///    = note: `-D clippy::forget-ref` implied by `-D warnings`
/// note: argument has type &SomeStruct
///   --> $DIR/drop_forget_ref.rs:10:12
///    |
/// 10 |     forget(&SomeStruct);
///    |            ^^^^^^^^^^^
/// ```
pub fn span_lint_and_note<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    span: Span,
    msg: &str,
    note_span: Option<Span>,
    note: &str,
) {
    cx.struct_span_lint(lint, span, |diag| {
        let mut diag = diag.build(msg);
        if let Some(note_span) = note_span {
            diag.span_note(note_span, note);
        } else {
            diag.note(note);
        }
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Like `span_lint` but allows to add notes, help and suggestions using a closure.
///
/// If you need to customize your lint output a lot, use this function.
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
pub fn span_lint_and_then<'a, T: LintContext, F>(cx: &'a T, lint: &'static Lint, sp: Span, msg: &str, f: F)
where
    F: for<'b> FnOnce(&mut DiagnosticBuilder<'b>),
{
    cx.struct_span_lint(lint, sp, |diag| {
        let mut diag = diag.build(msg);
        f(&mut diag);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

pub fn span_lint_hir(cx: &LateContext<'_>, lint: &'static Lint, hir_id: HirId, sp: Span, msg: &str) {
    cx.tcx.struct_span_lint_hir(lint, hir_id, sp, |diag| {
        let mut diag = diag.build(msg);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

pub fn span_lint_hir_and_then(
    cx: &LateContext<'_>,
    lint: &'static Lint,
    hir_id: HirId,
    sp: Span,
    msg: &str,
    f: impl FnOnce(&mut DiagnosticBuilder<'_>),
) {
    cx.tcx.struct_span_lint_hir(lint, hir_id, sp, |diag| {
        let mut diag = diag.build(msg);
        f(&mut diag);
        docs_link(&mut diag, lint);
        diag.emit();
    });
}

/// Add a span lint with a suggestion on how to fix it.
///
/// These suggestions can be parsed by rustfix to allow it to automatically fix your code.
/// In the example below, `help` is `"try"` and `sugg` is the suggested replacement `".any(|x| x >
/// 2)"`.
///
/// If you change the signature, remember to update the internal lint `CollapsibleCalls`
///
/// # Example
///
/// ```ignore
/// error: This `.fold` can be more succinctly expressed as `.any`
/// --> $DIR/methods.rs:390:13
///     |
/// 390 |     let _ = (0..3).fold(false, |acc, x| acc || x > 2);
///     |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: try: `.any(|x| x > 2)`
///     |
///     = note: `-D fold-any` implied by `-D warnings`
/// ```
#[cfg_attr(feature = "internal-lints", allow(clippy::collapsible_span_lint_calls))]
pub fn span_lint_and_sugg<'a, T: LintContext>(
    cx: &'a T,
    lint: &'static Lint,
    sp: Span,
    msg: &str,
    help: &str,
    sugg: String,
    applicability: Applicability,
) {
    span_lint_and_then(cx, lint, sp, msg, |diag| {
        diag.span_suggestion(sp, help, sugg, applicability);
    });
}

/// Create a suggestion made from several `span → replacement`.
///
/// Note: in the JSON format (used by `compiletest_rs`), the help message will
/// appear once per
/// replacement. In human-readable format though, it only appears once before
/// the whole suggestion.
pub fn multispan_sugg<I>(diag: &mut DiagnosticBuilder<'_>, help_msg: &str, sugg: I)
where
    I: IntoIterator<Item = (Span, String)>,
{
    multispan_sugg_with_applicability(diag, help_msg, Applicability::Unspecified, sugg)
}

pub fn multispan_sugg_with_applicability<I>(
    diag: &mut DiagnosticBuilder<'_>,
    help_msg: &str,
    applicability: Applicability,
    sugg: I,
) where
    I: IntoIterator<Item = (Span, String)>,
{
    diag.multipart_suggestion(help_msg, sugg.into_iter().collect(), applicability);
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

macro_rules! create_state_step {
    ($source:ident => $dest:ident) => {
        impl From<State<$source>> for State<$dest> {
            fn from(_from: State<$source>) -> State<$dest> {
                Self($dest {})
            }
        }
    };
    ($source:ident => $dest:ident($field:ident)) => {
        impl From<State<$source>> for State<$dest> {
            fn from(from: State<$source>) -> State<$dest> {
                Self($dest {
                    $field: (from.0).$field,
                })
            }
        }
    };
}

macro_rules! for_any_state {
    ($machine:ident, $state:ident, $code:block) => {
        match $machine {
            StateMachine::Park($state) => $code,
            StateMachine::Idle($state) => $code,
            StateMachine::Poll($state) => $code,
            StateMachine::Probe($state) => $code,
            StateMachine::Download($state) => $code,
            StateMachine::Install($state) => $code,
            StateMachine::Reboot($state) => $code,
        }
    };
}

macro_rules! shared_state {
    () => {
        crate::states::SHARED_STATE
            .clone()
            .read()
            .unwrap()
            .as_ref()
            .unwrap()
    };
}

macro_rules! shared_state_mut {
    () => {
        crate::states::SHARED_STATE
            .clone()
            .write()
            .unwrap()
            .as_mut()
            .unwrap()
    };
}

macro_rules! set_shared_state {
    ($settings:ident, $runtime_settings:ident, $firmware:ident) => {
        crate::states::SHARED_STATE
            .write()
            .unwrap()
            .replace(crate::states::SharedState {
                $settings,
                $runtime_settings,
                $firmware,
            });
    };
}

#[cfg(test)]
macro_rules! assert_state {
    ($machine:ident, $state:ident) => {
        assert!(
            if let Ok(StateMachine::$state(_)) = $machine {
                true
            } else {
                false
            },
            "Failed to get to {} state.",
            stringify!($state),
        );
    };
}
<commit_msg>updatehub: Change for_any_state macro to use full path for type<commit_after>// Copyright (C) 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

macro_rules! create_state_step {
    ($source:ident => $dest:ident) => {
        impl From<State<$source>> for State<$dest> {
            fn from(_from: State<$source>) -> State<$dest> {
                Self($dest {})
            }
        }
    };
    ($source:ident => $dest:ident($field:ident)) => {
        impl From<State<$source>> for State<$dest> {
            fn from(from: State<$source>) -> State<$dest> {
                Self($dest {
                    $field: (from.0).$field,
                })
            }
        }
    };
}

macro_rules! for_any_state {
    ($machine:ident, $state:ident, $code:block) => {
        match $machine {
            crate::states::StateMachine::Park($state) => $code,
            crate::states::StateMachine::Idle($state) => $code,
            crate::states::StateMachine::Poll($state) => $code,
            crate::states::StateMachine::Probe($state) => $code,
            crate::states::StateMachine::Download($state) => $code,
            crate::states::StateMachine::Install($state) => $code,
            crate::states::StateMachine::Reboot($state) => $code,
        }
    };
}

macro_rules! shared_state {
    () => {
        crate::states::SHARED_STATE
            .clone()
            .read()
            .unwrap()
            .as_ref()
            .unwrap()
    };
}

macro_rules! shared_state_mut {
    () => {
        crate::states::SHARED_STATE
            .clone()
            .write()
            .unwrap()
            .as_mut()
            .unwrap()
    };
}

macro_rules! set_shared_state {
    ($settings:ident, $runtime_settings:ident, $firmware:ident) => {
        crate::states::SHARED_STATE
            .write()
            .unwrap()
            .replace(crate::states::SharedState {
                $settings,
                $runtime_settings,
                $firmware,
            });
    };
}

#[cfg(test)]
macro_rules! assert_state {
    ($machine:ident, $state:ident) => {
        assert!(
            if let Ok(StateMachine::$state(_)) = $machine {
                true
            } else {
                false
            },
            "Failed to get to {} state.",
            stringify!($state),
        );
    };
}
<|endoftext|>


<commit_before>use std::fmt;
use hex_slice::AsHex;
use rand::{random, Rng, XorShiftRng};
use std::rc::Rc;
use std::cell::RefCell;

/// Something that can extract information from an `InfoSource`.
pub trait InfoSink {
    /// The output data.
    type Out;
    /// Called by [InfoSource::draw](trait.InfoSource.html#tymethod.draw)
    fn sink<I: InfoSource>(&mut self, i: &mut I) -> Self::Out;
}

/// Something that an act as a source of test data.
pub trait InfoSource {
    /// Take a single byte from the source.
    fn draw_u8(&mut self) -> u8;

    /// Call F with access to the data source.
    fn draw<S: InfoSink>(&mut self, sink: S) -> S::Out
    where
        Self: Sized;
}

/// Generates data from an underlying Rng instance.
#[derive(Debug)]
pub struct RngSource<R> {
    rng: R,
}

/// An adapter that can record the data drawn from an underlying source.
#[derive(Debug)]
pub struct InfoRecorder<I> {
    inner: Rc<RefCell<I>>,
    pub(crate) data: Vec<u8>,
}

impl<'a, I: InfoSource + ?Sized> InfoSource for &'a mut I {
    fn draw_u8(&mut self) -> u8 {
        (**self).draw_u8()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

impl<I> InfoRecorder<I> {
    /// Creates a recording InfoSource.
    pub fn new(inner: I) -> Self {
        InfoRecorder {
            inner: Rc::new(RefCell::new(inner)),
            data: Vec::new(),
        }
    }

    /// Extracts the data recorded.
    pub fn into_pool(self) -> InfoPool {
        InfoPool::of_vec(self.data)
    }
}

impl<I: InfoSource> InfoSource for InfoRecorder<I> {
    fn draw_u8(&mut self) -> u8 {
        let byte = self.inner.borrow_mut().draw_u8();
        self.data.push(byte);
        byte
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        trace!("-> InfoSource::draw");
        let res = sink.sink(self);
        trace!("<- InfoSource::draw");
        res
    }
}

impl RngSource<XorShiftRng> {
    /// Creates a RngSource with a randomly seeded XorShift generator.
    pub fn new() -> Self {
        let rng = random::<XorShiftRng>();
        RngSource { rng }
    }
}

impl<R: Rng> InfoSource for RngSource<R> {
    fn draw_u8(&mut self) -> u8 {
        self.rng.gen::<u8>()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

/// A pool of data that we can draw upon to generate other types of data.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InfoPool {
    pub(in data) data: Vec<u8>,
}

/// A handle to an info Pool that we can draw replayed bytes from, and zero after.
#[derive(Clone)]
pub struct InfoReplay<'a> {
    data: &'a [u8],
    off: usize,
}

impl fmt::Debug for InfoPool {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("InfoPool")
            .field("data", &format_args!("{:x}", self.data.as_hex()))
            .finish()
    }
}

/// The reasons why drawing data from a pool can fail.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum DataError {
    /// Not enough data to generate a value
    PoolExhausted,
    /// One of our combinators said that we should not test this value.
    SkipItem,
}

impl InfoPool {
    /// Create an `InfoPool` with a given vector of bytes. (Mostly used for
    /// testing).
    pub fn of_vec(data: Vec<u8>) -> Self {
        InfoPool { data: data }
    }

    /// Create an `InfoPool` with a `size` length vector of random bytes
    /// using the generator `rng`. (Mostly used for testing).
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Allows access to the underlying buffer.
    pub fn buffer(&self) -> &[u8] {
        &*self.data
    }

    /// Creates a tap that allows drawing new information from this pool.
    pub fn replay(&self) -> InfoReplay {
        InfoReplay {
            data: &*self.data,
            off: 0,
        }
    }
}

impl<'a> InfoReplay<'a> {
    /// Consumes the next byte from this tap. Returns `Ok(x)` if successful,
    /// or `Err(DataError::PoolExhausted)` if we have reached the end.
    pub fn next_byte(&mut self) -> u8 {
        let res = self.data.get(self.off).cloned();
        self.off += 1;
        res.unwrap_or(0)
    }
}

impl<'a> InfoSource for InfoReplay<'a> {
    fn draw_u8(&mut self) -> u8 {
        self.next_byte()
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use super::*;
    impl<R: Rng> RngSource<R> {
        pub(crate) fn of(rng: R) -> Self {
            RngSource { rng }
        }
    }

    #[test]
    fn should_take_each_item_in_pool() {
        let p = InfoPool::of_vec(vec![0, 1, 2, 3]);
        let mut t = p.replay();
        assert_eq!(t.next_byte(), 0);
        assert_eq!(t.next_byte(), 1);
        assert_eq!(t.next_byte(), 2);
        assert_eq!(t.next_byte(), 3);
        assert_eq!(t.next_byte(), 0);
    }

    #[test]
    fn should_generate_random_data() {
        let trials = 1024usize;
        let mut vals = 0;
        let mut p = RngSource::of(XorShiftRng::new_unseeded());
        let error = 8;
        for _ in 0..trials {
            vals += p.draw_u8() as usize;
        }
        let mean = vals / trials;
        let expected = 128;
        assert!(
            (expected - error) < mean && (expected + error) > mean,
            "Expected {} trials to be ({}+/-{}); got {}",
            trials,
            expected,
            error,
            mean
        )
    }

    #[test]
    fn should_allow_restarting_read() {
        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        {
            for _ in 0..4 {
                v0.push(p.draw_u8())
            }
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_mixed_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..8 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_borrowing_buffer() {
        let p = InfoPool::of_vec(vec![1]);
        assert_eq!(p.buffer(), &[1]);
    }

    #[test]
    fn tap_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());
        let _: &InfoSource = &p.replay();
        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

    #[test]
    fn replay_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());

        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

}
<commit_msg>Only increment offset in InfoReplay::next_byte iff we still have some left.<commit_after>use std::fmt;
use hex_slice::AsHex;
use rand::{random, Rng, XorShiftRng};
use std::rc::Rc;
use std::cell::RefCell;

/// Something that can extract information from an `InfoSource`.
pub trait InfoSink {
    /// The output data.
    type Out;
    /// Called by [InfoSource::draw](trait.InfoSource.html#tymethod.draw)
    fn sink<I: InfoSource>(&mut self, i: &mut I) -> Self::Out;
}

/// Something that an act as a source of test data.
pub trait InfoSource {
    /// Take a single byte from the source.
    fn draw_u8(&mut self) -> u8;

    /// Call F with access to the data source.
    fn draw<S: InfoSink>(&mut self, sink: S) -> S::Out
    where
        Self: Sized;
}

/// Generates data from an underlying Rng instance.
#[derive(Debug)]
pub struct RngSource<R> {
    rng: R,
}

/// An adapter that can record the data drawn from an underlying source.
#[derive(Debug)]
pub struct InfoRecorder<I> {
    inner: Rc<RefCell<I>>,
    pub(crate) data: Vec<u8>,
}

impl<'a, I: InfoSource + ?Sized> InfoSource for &'a mut I {
    fn draw_u8(&mut self) -> u8 {
        (**self).draw_u8()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

impl<I> InfoRecorder<I> {
    /// Creates a recording InfoSource.
    pub fn new(inner: I) -> Self {
        InfoRecorder {
            inner: Rc::new(RefCell::new(inner)),
            data: Vec::new(),
        }
    }

    /// Extracts the data recorded.
    pub fn into_pool(self) -> InfoPool {
        InfoPool::of_vec(self.data)
    }
}

impl<I: InfoSource> InfoSource for InfoRecorder<I> {
    fn draw_u8(&mut self) -> u8 {
        let byte = self.inner.borrow_mut().draw_u8();
        self.data.push(byte);
        byte
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        trace!("-> InfoSource::draw");
        let res = sink.sink(self);
        trace!("<- InfoSource::draw");
        res
    }
}

impl RngSource<XorShiftRng> {
    /// Creates a RngSource with a randomly seeded XorShift generator.
    pub fn new() -> Self {
        let rng = random::<XorShiftRng>();
        RngSource { rng }
    }
}

impl<R: Rng> InfoSource for RngSource<R> {
    fn draw_u8(&mut self) -> u8 {
        self.rng.gen::<u8>()
    }
    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

/// A pool of data that we can draw upon to generate other types of data.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InfoPool {
    pub(in data) data: Vec<u8>,
}

/// A handle to an info Pool that we can draw replayed bytes from, and zero after.
#[derive(Clone)]
pub struct InfoReplay<'a> {
    data: &'a [u8],
    off: usize,
}

impl fmt::Debug for InfoPool {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("InfoPool")
            .field("data", &format_args!("{:x}", self.data.as_hex()))
            .finish()
    }
}

/// The reasons why drawing data from a pool can fail.
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum DataError {
    /// Not enough data to generate a value
    PoolExhausted,
    /// One of our combinators said that we should not test this value.
    SkipItem,
}

impl InfoPool {
    /// Create an `InfoPool` with a given vector of bytes. (Mostly used for
    /// testing).
    pub fn of_vec(data: Vec<u8>) -> Self {
        InfoPool { data: data }
    }

    /// Create an `InfoPool` with a `size` length vector of random bytes
    /// using the generator `rng`. (Mostly used for testing).
    pub fn new() -> Self {
        Self { data: Vec::new() }
    }

    /// Allows access to the underlying buffer.
    pub fn buffer(&self) -> &[u8] {
        &*self.data
    }

    /// Creates a tap that allows drawing new information from this pool.
    pub fn replay(&self) -> InfoReplay {
        InfoReplay {
            data: &*self.data,
            off: 0,
        }
    }
}

impl<'a> InfoReplay<'a> {
    /// Consumes the next byte from this tap. Returns `Ok(x)` if successful,
    /// or `Err(DataError::PoolExhausted)` if we have reached the end.
    pub fn next_byte(&mut self) -> u8 {
        if let Some(res) = self.data.get(self.off).cloned() {
            self.off += 1;
            res
        } else {
            0
        }
    }
}

impl<'a> InfoSource for InfoReplay<'a> {
    fn draw_u8(&mut self) -> u8 {
        self.next_byte()
    }

    fn draw<S: InfoSink>(&mut self, mut sink: S) -> S::Out
    where
        Self: Sized,
    {
        sink.sink(self)
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use super::*;
    impl<R: Rng> RngSource<R> {
        pub(crate) fn of(rng: R) -> Self {
            RngSource { rng }
        }
    }

    #[test]
    fn should_take_each_item_in_pool() {
        let p = InfoPool::of_vec(vec![0, 1, 2, 3]);
        let mut t = p.replay();
        assert_eq!(t.next_byte(), 0);
        assert_eq!(t.next_byte(), 1);
        assert_eq!(t.next_byte(), 2);
        assert_eq!(t.next_byte(), 3);
        assert_eq!(t.next_byte(), 0);
    }

    #[test]
    fn should_generate_random_data() {
        let trials = 1024usize;
        let mut vals = 0;
        let mut p = RngSource::of(XorShiftRng::new_unseeded());
        let error = 8;
        for _ in 0..trials {
            vals += p.draw_u8() as usize;
        }
        let mean = vals / trials;
        let expected = 128;
        assert!(
            (expected - error) < mean && (expected + error) > mean,
            "Expected {} trials to be ({}+/-{}); got {}",
            trials,
            expected,
            error,
            mean
        )
    }

    #[test]
    fn should_allow_restarting_read() {
        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        {
            for _ in 0..4 {
                v0.push(p.draw_u8())
            }
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();
        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..4 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_restarting_mixed_child_reads() {
        struct FnSink<F>(F);
        impl<F: FnMut(&mut InfoSource) -> R, R> InfoSink for FnSink<F> {
            type Out = R;
            fn sink<I: InfoSource>(&mut self, k: &mut I) -> R {
                (self.0)(k as &mut InfoSource)
            }
        }

        let mut p = InfoRecorder::new(RngSource::new());
        let mut v0 = Vec::new();

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        p.draw(FnSink(|src: &mut InfoSource| {
            for _ in 0..4 {
                let x: u8 = src.draw_u8();
                v0.push(x);
            }
        }));

        for _ in 0..2 {
            let x: u8 = p.draw_u8();
            v0.push(x);
        }

        let p = p.into_pool();
        let mut t = p.replay();
        let mut v1 = Vec::new();
        for _ in 0..8 {
            v1.push(t.draw_u8())
        }

        assert_eq!(v0, v1)
    }

    #[test]
    fn should_allow_borrowing_buffer() {
        let p = InfoPool::of_vec(vec![1]);
        assert_eq!(p.buffer(), &[1]);
    }

    #[test]
    fn tap_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());
        let _: &InfoSource = &p.replay();
        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

    #[test]
    fn replay_can_act_as_source() {
        let buf = vec![4, 3, 2, 1];
        let p = InfoPool::of_vec(buf.clone());

        let mut res = Vec::new();
        let mut it = p.replay();
        for _ in 0..4 {
            res.push(it.draw_u8())
        }
        assert_eq!(res, buf)
    }

}
<|endoftext|>


<commit_before>extern crate alfred;

use config::Config;
use jira_api::SearchIssue;
use jira_client;
use search_response::SearchResponse;
use std::io::{self};

pub struct SearchCommand {
    pub config: Config,
}

impl SearchCommand {
    pub fn run(&self, keyword: &String) -> Result<(), String> {
        let jira_search = SearchIssue{
            config: self.config.clone(),
            keyword: keyword.to_string(),
        };
        let res = jira_client::send::<SearchIssue, SearchResponse>(jira_search).map_err(|e| e.to_string())?;
        Ok(self.write_alfred_items(res))
    }

    fn write_alfred_items(&self, res: SearchResponse) {
        let alfred_items = res.issues.into_iter()
            .map(|issue| alfred::ItemBuilder::new(issue.key).arg(issue.url).into_item())
            .collect::<Vec<alfred::Item>>();

        alfred::json::write_items(io::stdout(), &alfred_items);
    }
}<commit_msg>:art: Make possible to throw Err when write alfred items<commit_after>extern crate alfred;

use config::Config;
use jira_api::SearchIssue;
use jira_client;
use search_response::SearchResponse;
use std::io::{self};

pub struct SearchCommand {
    pub config: Config,
}

impl SearchCommand {
    pub fn run(&self, keyword: &String) -> Result<(), String> {
        let jira_search = SearchIssue {
            config: self.config.clone(),
            keyword: keyword.to_string(),
        };
        let res = jira_client::send::<SearchIssue, SearchResponse>(jira_search).map_err(|e| e.to_string())?;
        Ok(self.write_alfred_items(res)?)
    }

    fn write_alfred_items(&self, res: SearchResponse) -> Result<(), String> {
        let alfred_items = res.issues.into_iter()
            .map(|issue| alfred::ItemBuilder::new(issue.key).arg(issue.url).into_item())
            .collect::<Vec<alfred::Item>>();

        alfred::json::write_items(io::stdout(), &alfred_items).map_err(|e| e.to_string())
    }
}<|endoftext|>


<commit_before>/*!
 * Just make sure that the regular expression accepts the sorts of things we want to
 * accept. Don't test what this grammar parses just yet.
 */

rusty_peg! {
    parser Parser<'input> {
        StrHi: &'input str = "Hi";
        StrHi1: &'input str = ("Hi");
        StrHiHo: (&'input str, &'input str) = ("Hi" "Ho");

        IntHi: u32 = "Hi" => 1;
        IntHi1: u32 = ("Hi") => 1;
        IntHiHo: (u32, u32) = ("Hi" "Ho") => (1,1);
    }
}

<commit_msg>Add some more combos.<commit_after>/*!
 * Just make sure that the regular expression accepts the sorts of things we want to
 * accept. Don't test what this grammar parses just yet.
 */

rusty_peg! {
    parser Parser<'input> {
        StrHi: &'input str = "Hi";
        StrHi1: &'input str = ("Hi");
        StrHiHo: (&'input str, &'input str) = ("Hi" "Ho");

        IntHi: u32 = "Hi" => 1;
        IntHi1: u32 = ("Hi") => 1;
        IntHiHo: (u32, u32) = ("Hi" "Ho") => (1,1);

        VecHi: Vec<&'input str> = {"Hi"};
        VecHiHo: Vec<(&'input str, &'input str)> = {"Hi" "Ho"};

        OptHi: Option<&'input str> = ["Hi"];
        OptHiHo: Option<(&'input str, &'input str)> = ["Hi" "Ho"];
    }
}

<|endoftext|>


<commit_before>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {:#x}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 144)?,
            false => ws.read_u32(load_config_directory + 88)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 128)?,
            false => ws.read_va(load_config_directory + 80)?,
        };

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 136)?,
            false => ws.read_u32(load_config_directory + 84)?,
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        let mut offset = cfg_table_rva;

        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {:#x}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 112)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 72)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 120)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 76)?).unwrap(),
        };
        if ws.probe(guard_dispatch_icall_fptr, 8) {
            let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
            if ws.probe(guard_dispatch_icall, 1) {
                debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                ws.make_function(guard_dispatch_icall)?;
                ws.analyze()?;
            }
        };

        Ok(())
    }
}
<commit_msg>*: dbg formatting<commit_after>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 144)?,
            false => ws.read_u32(load_config_directory + 88)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 128)?,
            false => ws.read_va(load_config_directory + 80)?,
        };

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 136)?,
            false => ws.read_u32(load_config_directory + 84)?,
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        let mut offset = cfg_table_rva;

        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 112)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 72)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 120)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 76)?).unwrap(),
        };
        if ws.probe(guard_dispatch_icall_fptr, 8) {
            let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
            if ws.probe(guard_dispatch_icall, 1) {
                debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                ws.make_function(guard_dispatch_icall)?;
                ws.analyze()?;
            }
        };

        Ok(())
    }
}
<|endoftext|>