<commit_before>use schema::*;
use diesel::*;
use diesel::backend::Pg;

#[test]
fn filter_by_int_equality() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean_id = find_user_by_name("Sean", &connection).id;
    let tess_id = find_user_by_name("Tess", &connection).id;
    let unused_id = sean_id + tess_id;

    let sean = User::new(sean_id, "Sean");
    let tess = User::new(tess_id, "Tess");
    assert_eq!(Ok(sean), users.filter(id.eq(sean_id)).first(&connection));
    assert_eq!(Ok(tess), users.filter(id.eq(tess_id)).first(&connection));
    assert_eq!(Err(NotFound), users.filter(id.eq(unused_id)).first::<User>(&connection));
}

#[test]
fn filter_by_string_equality() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();

    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    assert_eq!(Ok(sean), users.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok(tess), users.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound), users.filter(name.eq("Jim")).first::<User>(&connection));
}

#[test]
fn sqlite_empty() {
    use diesel::connection::SqliteConnection;
    use schema::users::dsl::*;

    let connection = SqliteConnection::establish(":memory:").unwrap();
    connection.execute("
        CREATE TABLE users (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          name VARCHAR NOT NULL,
          hair_color VARCHAR
        )").unwrap();

    assert_eq!(Vec::<User>::new(), users.load(&connection).unwrap().collect::<Vec<User>>());
}

#[test]
fn zomg_sqlite() {
    use diesel::connection::SqliteConnection;
    use schema::users::dsl::*;

    let connection = SqliteConnection::establish(":memory:").unwrap();
    connection.execute("
        CREATE TABLE users (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          name VARCHAR NOT NULL,
          hair_color VARCHAR
        )").unwrap();
    insert(&vec![NewUser::new("Sean", None), NewUser::new("Tess", None)])
        .into(users)
        .execute(&connection)
        .unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let sean = &data[0];
    let tess = &data[1];

    assert_eq!(Ok(sean), users.filter(name.eq("Sean")).first(&connection).as_ref());
    assert_eq!(Ok(tess), users.filter(name.eq("Tess")).first(&connection).as_ref());
    assert_eq!(Err(NotFound), users.filter(name.eq("Jim")).first::<User>(&connection));
}

#[test]
fn filter_by_equality_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Sean", Some("black")),
        NewUser::new("Tess", Some("brown")),
        NewUser::new("Jim", Some("black")),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let sean = data[0].clone();
    let tess = data[1].clone();
    let jim = data[2].clone();

    let source = users.filter(hair_color.eq("black"));
    assert_eq!(vec![sean, jim], source.load(&connection).unwrap().collect::<Vec<_>>());
    let source = users.filter(hair_color.eq("brown"));
    assert_eq!(vec![tess], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
fn filter_by_is_not_null_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Derek", Some("red")),
        NewUser::new("Gordon", None),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let derek = data[0].clone();

    let source = users.filter(hair_color.is_not_null());
    assert_eq!(vec![derek], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
fn filter_by_is_null_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Derek", Some("red")),
        NewUser::new("Gordon", None),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let gordon = data[1].clone();

    let source = users.filter(hair_color.is_null());
    assert_eq!(vec![gordon], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn filter_after_joining() {
    use schema::users::name;

    let connection = connection_with_sean_and_tess_in_users_table();
    connection.execute("INSERT INTO POSTS (id, title, user_id) VALUES
                       (1, 'Hello', 1), (2, 'World', 2)")
        .unwrap();

    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    let seans_post = Post::new(1, 1, "Hello", None);
    let tess_post = Post::new(2, 2, "World", None);
    let source = users::table.inner_join(posts::table);
    assert_eq!(Ok((sean, seans_post)),
        source.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok((tess, tess_post)),
        source.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound),
        source.filter(name.eq("Jim")).first::<(User, Post)>(&connection));
}

#[test]
fn select_then_filter() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();

    let source = users.select(name);
    assert_eq!(Ok("Sean".to_string()),
        source.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok("Tess".to_string()),
        source.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound), source.filter(name.eq("Jim")).first::<String>(&connection));
}

#[test]
fn filter_then_select() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![NewUser::new("Sean", None), NewUser::new("Tess", None)];
    insert(&data).into(users).execute(&connection).unwrap();

    assert_eq!(Ok("Sean".to_string()),
        users.filter(name.eq("Sean")).select(name).first(&connection));
    assert_eq!(Ok("Tess".to_string()),
        users.filter(name.eq("Tess")).select(name).first(&connection));
    assert_eq!(Err(NotFound), users.filter(name.eq("Jim")).select(name)
                                   .first::<String>(&connection));
}

#[test]
fn filter_on_multiple_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data: &[_] = &[
        NewUser::new("Sean", Some("black")),
        NewUser::new("Sean", Some("brown")),
        NewUser::new("Sean", None),
        NewUser::new("Tess", Some("black")),
        NewUser::new("Tess", Some("brown")),
    ];
    insert(data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let black_haired_sean = data[0].clone();
    let brown_haired_sean = data[1].clone();
    let black_haired_tess = data[3].clone();
    let brown_haired_tess = data[4].clone();

    let source = users.filter(name.eq("Sean").and(hair_color.eq("black")));
    assert_eq!(vec![black_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Sean").and(hair_color.eq("brown")));
    assert_eq!(vec![brown_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess").and(hair_color.eq("black")));
    assert_eq!(vec![black_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess").and(hair_color.eq("brown")));
    assert_eq!(vec![brown_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());
}

#[test]
fn filter_called_twice_means_same_thing_as_and() {
    use schema::users::dsl::*;

    let connection = connection();
    let data: &[_] = &[
        NewUser::new("Sean", Some("black")),
        NewUser::new("Sean", Some("brown")),
        NewUser::new("Sean", None),
        NewUser::new("Tess", Some("black")),
        NewUser::new("Tess", Some("brown")),
    ];
    insert(data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let black_haired_sean = data[0].clone();
    let brown_haired_sean = data[1].clone();
    let black_haired_tess = data[3].clone();
    let brown_haired_tess = data[4].clone();

    let source = users.filter(name.eq("Sean")).filter(hair_color.eq("black"));
    assert_eq!(vec![black_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Sean")).filter(hair_color.eq("brown"));
    assert_eq!(vec![brown_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess")).filter(hair_color.eq("black"));
    assert_eq!(vec![black_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess")).filter(hair_color.eq("brown"));
    assert_eq!(vec![brown_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());
}

table! {
    points (x) {
        x -> Integer,
        y -> Integer,
    }
}

#[test]
fn filter_on_column_equality() {
    use self::points::dsl::*;

    let connection = connection();
    connection.execute("CREATE TABLE points (x INTEGER NOT NULL, y INTEGER NOT NULL)").unwrap();
    connection.execute("INSERT INTO POINTS (x, y) VALUES (1, 1), (1, 2), (2, 2)").unwrap();

    let expected_data = vec![(1, 1), (2, 2)];
    let query = points.filter(x.eq(y));
    let data: Vec<_> = query.load(&connection).unwrap().collect();
    assert_eq!(expected_data, data);
}

#[test]
fn filter_with_or() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    insert(&NewUser::new("Jim", None)).into(users).execute(&connection).unwrap();

    let expected_users = vec![User::new(1, "Sean"), User::new(2, "Tess")];
    let data: Vec<_> = users.filter(name.eq("Sean").or(name.eq("Tess")))
        .load(&connection).unwrap().collect();

    assert_eq!(expected_users, data);
}

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn or_doesnt_mess_with_precidence_of_previous_statements() {
    use schema::users::dsl::*;
    use diesel::expression::AsExpression;

    let connection = connection_with_sean_and_tess_in_users_table();
    let f = AsExpression::<types::Bool>::as_expression(false);
    let count = users.filter(f).filter(f.or(true))
        .count().first(&connection);

    assert_eq!(Ok(0), count);

    let count = users.filter(f.or(f).and(f.or(true)))
        .count().first(&connection);

    assert_eq!(Ok(0), count);
}

use diesel::types::VarChar;
sql_function!(lower, lower_t, (x: VarChar) -> VarChar);

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn filter_by_boxed_predicate() {
    fn by_name(name: &str) -> Box<BoxableExpression<users::table, types::Bool, Pg, SqlType=types::Bool>> {
        Box::new(lower(users::name).eq(name.to_string()))
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    let queried_sean = users::table.filter(by_name("sean")).first(&connection);
    let queried_tess = users::table.filter(by_name("tess")).first(&connection);

    assert_eq!(Ok(sean), queried_sean);
    assert_eq!(Ok(tess), queried_tess);
}
<commit_msg>Remove bad tests<commit_after>use schema::*;
use diesel::*;
use diesel::backend::Pg;

#[test]
fn filter_by_int_equality() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean_id = find_user_by_name("Sean", &connection).id;
    let tess_id = find_user_by_name("Tess", &connection).id;
    let unused_id = sean_id + tess_id;

    let sean = User::new(sean_id, "Sean");
    let tess = User::new(tess_id, "Tess");
    assert_eq!(Ok(sean), users.filter(id.eq(sean_id)).first(&connection));
    assert_eq!(Ok(tess), users.filter(id.eq(tess_id)).first(&connection));
    assert_eq!(Err(NotFound), users.filter(id.eq(unused_id)).first::<User>(&connection));
}

#[test]
fn filter_by_string_equality() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();

    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    assert_eq!(Ok(sean), users.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok(tess), users.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound), users.filter(name.eq("Jim")).first::<User>(&connection));
}

#[test]
fn filter_by_equality_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Sean", Some("black")),
        NewUser::new("Tess", Some("brown")),
        NewUser::new("Jim", Some("black")),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let sean = data[0].clone();
    let tess = data[1].clone();
    let jim = data[2].clone();

    let source = users.filter(hair_color.eq("black"));
    assert_eq!(vec![sean, jim], source.load(&connection).unwrap().collect::<Vec<_>>());
    let source = users.filter(hair_color.eq("brown"));
    assert_eq!(vec![tess], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
fn filter_by_is_not_null_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Derek", Some("red")),
        NewUser::new("Gordon", None),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let derek = data[0].clone();

    let source = users.filter(hair_color.is_not_null());
    assert_eq!(vec![derek], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
fn filter_by_is_null_on_nullable_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![
        NewUser::new("Derek", Some("red")),
        NewUser::new("Gordon", None),
    ];
    insert(&data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let gordon = data[1].clone();

    let source = users.filter(hair_color.is_null());
    assert_eq!(vec![gordon], source.load(&connection).unwrap().collect::<Vec<_>>());
}

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn filter_after_joining() {
    use schema::users::name;

    let connection = connection_with_sean_and_tess_in_users_table();
    connection.execute("INSERT INTO POSTS (id, title, user_id) VALUES
                       (1, 'Hello', 1), (2, 'World', 2)")
        .unwrap();

    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    let seans_post = Post::new(1, 1, "Hello", None);
    let tess_post = Post::new(2, 2, "World", None);
    let source = users::table.inner_join(posts::table);
    assert_eq!(Ok((sean, seans_post)),
        source.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok((tess, tess_post)),
        source.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound),
        source.filter(name.eq("Jim")).first::<(User, Post)>(&connection));
}

#[test]
fn select_then_filter() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();

    let source = users.select(name);
    assert_eq!(Ok("Sean".to_string()),
        source.filter(name.eq("Sean")).first(&connection));
    assert_eq!(Ok("Tess".to_string()),
        source.filter(name.eq("Tess")).first(&connection));
    assert_eq!(Err(NotFound), source.filter(name.eq("Jim")).first::<String>(&connection));
}

#[test]
fn filter_then_select() {
    use schema::users::dsl::*;

    let connection = connection();
    let data = vec![NewUser::new("Sean", None), NewUser::new("Tess", None)];
    insert(&data).into(users).execute(&connection).unwrap();

    assert_eq!(Ok("Sean".to_string()),
        users.filter(name.eq("Sean")).select(name).first(&connection));
    assert_eq!(Ok("Tess".to_string()),
        users.filter(name.eq("Tess")).select(name).first(&connection));
    assert_eq!(Err(NotFound), users.filter(name.eq("Jim")).select(name)
                                   .first::<String>(&connection));
}

#[test]
fn filter_on_multiple_columns() {
    use schema::users::dsl::*;

    let connection = connection();
    let data: &[_] = &[
        NewUser::new("Sean", Some("black")),
        NewUser::new("Sean", Some("brown")),
        NewUser::new("Sean", None),
        NewUser::new("Tess", Some("black")),
        NewUser::new("Tess", Some("brown")),
    ];
    insert(data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let black_haired_sean = data[0].clone();
    let brown_haired_sean = data[1].clone();
    let black_haired_tess = data[3].clone();
    let brown_haired_tess = data[4].clone();

    let source = users.filter(name.eq("Sean").and(hair_color.eq("black")));
    assert_eq!(vec![black_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Sean").and(hair_color.eq("brown")));
    assert_eq!(vec![brown_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess").and(hair_color.eq("black")));
    assert_eq!(vec![black_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess").and(hair_color.eq("brown")));
    assert_eq!(vec![brown_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());
}

#[test]
fn filter_called_twice_means_same_thing_as_and() {
    use schema::users::dsl::*;

    let connection = connection();
    let data: &[_] = &[
        NewUser::new("Sean", Some("black")),
        NewUser::new("Sean", Some("brown")),
        NewUser::new("Sean", None),
        NewUser::new("Tess", Some("black")),
        NewUser::new("Tess", Some("brown")),
    ];
    insert(data).into(users).execute(&connection).unwrap();
    let data = users.load(&connection).unwrap().collect::<Vec<User>>();
    let black_haired_sean = data[0].clone();
    let brown_haired_sean = data[1].clone();
    let black_haired_tess = data[3].clone();
    let brown_haired_tess = data[4].clone();

    let source = users.filter(name.eq("Sean")).filter(hair_color.eq("black"));
    assert_eq!(vec![black_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Sean")).filter(hair_color.eq("brown"));
    assert_eq!(vec![brown_haired_sean], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess")).filter(hair_color.eq("black"));
    assert_eq!(vec![black_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());

    let source = users.filter(name.eq("Tess")).filter(hair_color.eq("brown"));
    assert_eq!(vec![brown_haired_tess], source.load(&connection).unwrap()
        .collect::<Vec<_>>());
}

table! {
    points (x) {
        x -> Integer,
        y -> Integer,
    }
}

#[test]
fn filter_on_column_equality() {
    use self::points::dsl::*;

    let connection = connection();
    connection.execute("CREATE TABLE points (x INTEGER NOT NULL, y INTEGER NOT NULL)").unwrap();
    connection.execute("INSERT INTO POINTS (x, y) VALUES (1, 1), (1, 2), (2, 2)").unwrap();

    let expected_data = vec![(1, 1), (2, 2)];
    let query = points.filter(x.eq(y));
    let data: Vec<_> = query.load(&connection).unwrap().collect();
    assert_eq!(expected_data, data);
}

#[test]
fn filter_with_or() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    insert(&NewUser::new("Jim", None)).into(users).execute(&connection).unwrap();

    let expected_users = vec![User::new(1, "Sean"), User::new(2, "Tess")];
    let data: Vec<_> = users.filter(name.eq("Sean").or(name.eq("Tess")))
        .load(&connection).unwrap().collect();

    assert_eq!(expected_users, data);
}

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn or_doesnt_mess_with_precidence_of_previous_statements() {
    use schema::users::dsl::*;
    use diesel::expression::AsExpression;

    let connection = connection_with_sean_and_tess_in_users_table();
    let f = AsExpression::<types::Bool>::as_expression(false);
    let count = users.filter(f).filter(f.or(true))
        .count().first(&connection);

    assert_eq!(Ok(0), count);

    let count = users.filter(f.or(f).and(f.or(true)))
        .count().first(&connection);

    assert_eq!(Ok(0), count);
}

use diesel::types::VarChar;
sql_function!(lower, lower_t, (x: VarChar) -> VarChar);

#[test]
#[cfg(feature = "postgres")] // FIXME: There are valuable tests for SQLite here
fn filter_by_boxed_predicate() {
    fn by_name(name: &str) -> Box<BoxableExpression<users::table, types::Bool, Pg, SqlType=types::Bool>> {
        Box::new(lower(users::name).eq(name.to_string()))
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = User::new(1, "Sean");
    let tess = User::new(2, "Tess");
    let queried_sean = users::table.filter(by_name("sean")).first(&connection);
    let queried_tess = users::table.filter(by_name("tess")).first(&connection);

    assert_eq!(Ok(sean), queried_sean);
    assert_eq!(Ok(tess), queried_tess);
}
<|endoftext|>


<commit_before>use diesel::*;
use schema::*;

#[test]
fn test_updating_single_column() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    update(users)
        .set(name.eq("Jim"))
        .execute(&connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(); 2];
    let data: Vec<String> = users.select(name).load(&connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_single_column_of_single_row() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set(name.eq("Jim"))
        .execute(&connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(), "Tess".to_string()];
    let data: Vec<String> = users.select(name).order(id).load(&connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_nullable_column() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(Some("black")))
        .execute(&connection)
        .unwrap();

    let data: Option<String> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(&connection)
        .unwrap();
    assert_eq!(Some("black".to_string()), data);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(None::<String>))
        .execute(&connection)
        .unwrap();

    let data: QueryResult<Option<String>> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(&connection);
    assert_eq!(Ok(None::<String>), data);
}

#[test]
fn test_updating_multiple_columns() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set((name.eq("Jim"), hair_color.eq(Some("black"))))
        .execute(&connection)
        .unwrap();

    let expected_user = User::with_hair_color(sean.id, "Jim", "black");
    let user = users.find(sean.id).first(&connection);
    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(feature = "sqlite", feature = "mysql")))]
fn update_returning_struct() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .get_result(&connection);
    let expected_user = User::with_hair_color(sean.id, "Sean", "black");

    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(feature = "sqlite", feature = "mysql")))]
fn update_with_custom_returning_clause() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .returning((name, hair_color))
        .get_result::<(String, Option<String>)>(&connection);
    let expected_result = ("Sean".to_string(), Some("black".to_string()));

    assert_eq!(Ok(expected_result), user);
}

#[test]
fn update_with_struct_as_changes() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let changes = NewUser::new("Jim", Some("blue"));

    update(users.filter(id.eq(sean.id)))
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let user = users.find(sean.id).first(&connection);
    let expected_user = User::with_hair_color(sean.id, "Jim", "blue");

    assert_eq!(Ok(expected_user), user);
}

#[test]
fn save_on_struct_with_primary_key_changes_that_struct() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = User::with_hair_color(sean.id, "Jim", "blue").save_changes::<User>(&connection);

    let user_in_db = users.find(sean.id).first(&connection);

    assert_eq!(user, user_in_db);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_before_non_option() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        hair_color: Option<String>,
        name: String,
    }

    let changes = Changes {
        hair_color: None,
        name: "Jim".into(),
    };
    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    update(users::table.filter(users::id.eq(sean.id)))
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let user = users::table.find(sean.id).first(&connection);

    let expected_user = User::new(sean.id, "Jim");
    assert_eq!(Ok(expected_user), user);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_mixed_with_non_option() {
    #[derive(AsChangeset)]
    #[table_name = "posts"]
    struct Changes {
        user_id: i32,
        title: Option<String>,
        body: String,
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let new_post = sean.new_post("Hello", Some("world"));
    insert_into(posts::table)
        .values(&new_post)
        .execute(&connection)
        .unwrap();

    let changes = Changes {
        user_id: 1,
        title: None,
        body: "earth".into(),
    };
    update(posts::table)
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let post = posts::table
        .order(posts::id.desc())
        .first::<Post>(&connection)
        .unwrap();

    let expected_post = Post::new(post.id, sean.id, "Hello".into(), Some("earth".into()));
    assert_eq!(expected_post, post);
}

#[test]
#[should_panic(expected = "There are no changes to save.")]
fn update_with_no_changes() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        name: Option<String>,
        hair_color: Option<String>,
    }

    let connection = connection();
    let changes = Changes {
        name: None,
        hair_color: None,
    };
    update(users::table)
        .set(&changes)
        .execute(&connection)
        .unwrap();
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_no_changes_executes_do_nothing() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(&User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(&connection);

    assert_eq!(Ok(0), result);

    // Try the same thing with an owned type.
    let connection = connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(&connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_sql_literal_for_target() {
    use diesel::dsl::sql;
    use diesel::pg::upsert::*;
    use diesel::sql_types::Text;
    use schema::users::dsl::*;

    let connection = connection();
    // This index needs to happen before the insert or we'll get a deadlock
    // with any transactions that are trying to get the row lock from insert
    connection
        .execute("CREATE UNIQUE INDEX ON users (name) WHERE name != 'Tess'")
        .unwrap();
    insert_sean_and_tess_into_users_table(&connection);

    let new_users = vec![
        NewUser::new("Sean", Some("Green")),
        NewUser::new("Tess", Some("Blue")),
    ];
    insert_into(users)
        .values(&new_users)
        .on_conflict(sql::<Text>("(name) WHERE name != 'Tess'"))
        .do_update()
        .set(hair_color.eq(excluded(hair_color)))
        .execute(&connection)
        .unwrap();

    let data = users.select((name, hair_color)).order(id).load(&connection);
    let expected_data = vec![
        ("Sean".to_string(), Some("Green".to_string())),
        ("Tess".to_string(), None),
        ("Tess".to_string(), Some("Blue".to_string())),
    ];
    assert_eq!(Ok(expected_data), data);
}
<commit_msg>Fix test which was hanging forever<commit_after>use diesel::*;
use schema::*;

#[test]
fn test_updating_single_column() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    update(users)
        .set(name.eq("Jim"))
        .execute(&connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(); 2];
    let data: Vec<String> = users.select(name).load(&connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_single_column_of_single_row() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set(name.eq("Jim"))
        .execute(&connection)
        .unwrap();

    let expected_data = vec!["Jim".to_string(), "Tess".to_string()];
    let data: Vec<String> = users.select(name).order(id).load(&connection).unwrap();
    assert_eq!(expected_data, data);
}

#[test]
fn test_updating_nullable_column() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(Some("black")))
        .execute(&connection)
        .unwrap();

    let data: Option<String> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(&connection)
        .unwrap();
    assert_eq!(Some("black".to_string()), data);

    update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq(None::<String>))
        .execute(&connection)
        .unwrap();

    let data: QueryResult<Option<String>> = users
        .select(hair_color)
        .filter(id.eq(sean.id))
        .first(&connection);
    assert_eq!(Ok(None::<String>), data);
}

#[test]
fn test_updating_multiple_columns() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);

    update(users.filter(id.eq(sean.id)))
        .set((name.eq("Jim"), hair_color.eq(Some("black"))))
        .execute(&connection)
        .unwrap();

    let expected_user = User::with_hair_color(sean.id, "Jim", "black");
    let user = users.find(sean.id).first(&connection);
    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(feature = "sqlite", feature = "mysql")))]
fn update_returning_struct() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .get_result(&connection);
    let expected_user = User::with_hair_color(sean.id, "Sean", "black");

    assert_eq!(Ok(expected_user), user);
}

#[test]
#[cfg(not(any(feature = "sqlite", feature = "mysql")))]
fn update_with_custom_returning_clause() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = update(users.filter(id.eq(sean.id)))
        .set(hair_color.eq("black"))
        .returning((name, hair_color))
        .get_result::<(String, Option<String>)>(&connection);
    let expected_result = ("Sean".to_string(), Some("black".to_string()));

    assert_eq!(Ok(expected_result), user);
}

#[test]
fn update_with_struct_as_changes() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let changes = NewUser::new("Jim", Some("blue"));

    update(users.filter(id.eq(sean.id)))
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let user = users.find(sean.id).first(&connection);
    let expected_user = User::with_hair_color(sean.id, "Jim", "blue");

    assert_eq!(Ok(expected_user), user);
}

#[test]
fn save_on_struct_with_primary_key_changes_that_struct() {
    use schema::users::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let user = User::with_hair_color(sean.id, "Jim", "blue").save_changes::<User>(&connection);

    let user_in_db = users.find(sean.id).first(&connection);

    assert_eq!(user, user_in_db);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_before_non_option() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        hair_color: Option<String>,
        name: String,
    }

    let changes = Changes {
        hair_color: None,
        name: "Jim".into(),
    };
    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    update(users::table.filter(users::id.eq(sean.id)))
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let user = users::table.find(sean.id).first(&connection);

    let expected_user = User::new(sean.id, "Jim");
    assert_eq!(Ok(expected_user), user);
}

#[test]
fn sql_syntax_is_correct_when_option_field_comes_mixed_with_non_option() {
    #[derive(AsChangeset)]
    #[table_name = "posts"]
    struct Changes {
        user_id: i32,
        title: Option<String>,
        body: String,
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let new_post = sean.new_post("Hello", Some("world"));
    insert_into(posts::table)
        .values(&new_post)
        .execute(&connection)
        .unwrap();

    let changes = Changes {
        user_id: 1,
        title: None,
        body: "earth".into(),
    };
    update(posts::table)
        .set(&changes)
        .execute(&connection)
        .unwrap();
    let post = posts::table
        .order(posts::id.desc())
        .first::<Post>(&connection)
        .unwrap();

    let expected_post = Post::new(post.id, sean.id, "Hello".into(), Some("earth".into()));
    assert_eq!(expected_post, post);
}

#[test]
#[should_panic(expected = "There are no changes to save.")]
fn update_with_no_changes() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        name: Option<String>,
        hair_color: Option<String>,
    }

    let connection = connection();
    let changes = Changes {
        name: None,
        hair_color: None,
    };
    update(users::table)
        .set(&changes)
        .execute(&connection)
        .unwrap();
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_no_changes_executes_do_nothing() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(&User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(&connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_no_changes_executes_do_nothing_owned() {
    #[derive(AsChangeset)]
    #[table_name = "users"]
    struct Changes {
        hair_color: Option<String>,
    }

    let connection = connection_with_sean_and_tess_in_users_table();
    let result = insert_into(users::table)
        .values(User::new(1, "Sean"))
        .on_conflict(users::id)
        .do_update()
        .set(&Changes { hair_color: None })
        .execute(&connection);

    assert_eq!(Ok(0), result);
}

#[test]
#[cfg(feature = "postgres")]
fn upsert_with_sql_literal_for_target() {
    use diesel::dsl::sql;
    use diesel::pg::upsert::*;
    use diesel::sql_types::Text;
    use schema::users::dsl::*;

    let connection = connection();
    // This index needs to happen before the insert or we'll get a deadlock
    // with any transactions that are trying to get the row lock from insert
    connection
        .execute("CREATE UNIQUE INDEX ON users (name) WHERE name != 'Tess'")
        .unwrap();
    insert_sean_and_tess_into_users_table(&connection);

    let new_users = vec![
        NewUser::new("Sean", Some("Green")),
        NewUser::new("Tess", Some("Blue")),
    ];
    insert_into(users)
        .values(&new_users)
        .on_conflict(sql::<Text>("(name) WHERE name != 'Tess'"))
        .do_update()
        .set(hair_color.eq(excluded(hair_color)))
        .execute(&connection)
        .unwrap();

    let data = users.select((name, hair_color)).order(id).load(&connection);
    let expected_data = vec![
        ("Sean".to_string(), Some("Green".to_string())),
        ("Tess".to_string(), None),
        ("Tess".to_string(), Some("Blue".to_string())),
    ];
    assert_eq!(Ok(expected_data), data);
}
<|endoftext|>


<commit_before>/*
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

extern crate aplib;
extern crate docopt;
extern crate rustc_serialize;
extern crate pbr;

use docopt::Docopt;
use rustc_serialize::{Decodable, Decoder};
use pbr::ProgressBar;

use aplib::AplibObject;
use aplib::Library;
use aplib::Keyword;
use aplib::StoreWrapper;
use aplib::audit::Reporter;
use aplib::AlbumSubclass;
use aplib::FolderType;

const USAGE: &'static str = "
Usage:
  dumper <command> ([--all] | [--albums] [--versions] [--masters] [--folders] [--keywords]) <path>

Options:
    --all          Select all objects
    --albums       Select only albums
    --masters      Select only masters
    --folders      Select only folders
    --keywords     Select only keywords

Commands are:
    dump           Dump the objects
    audit          Audit mode: output what we ignored
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_command: Command,
    arg_path: String,
    flag_all: bool,
    flag_albums: bool,
    flag_versions: bool,
    flag_masters: bool,
    flag_folders: bool,
    flag_keywords: bool
}

#[derive(Debug)]
enum Command {
    Dump,
    Audit,
    Unknown(String)
}

impl Decodable for Command {
    fn decode<D: Decoder>(d: &mut D) -> Result<Command, D::Error> {
        let s = try!(d.read_str());
        Ok(match &*s {
            "dump" => Command::Dump,
            "audit" => Command::Audit,
            s => Command::Unknown(s.to_owned()),
        })
    }
}

fn main() {

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    {
        match args.arg_command {
            Command::Dump =>
                process_dump(&args),
            Command::Audit =>
                process_audit(&args),
            _ =>
                ()
        };
    }
}

fn process_audit(args: &Args) {

    let mut library = Library::new(&args.arg_path);

    let auditor = Reporter::new();
    library.set_auditor(Some(auditor));

    library.library_version();
    library.load_folders(&mut |_: u64| true);
    library.load_albums(&mut |_: u64| true);
    library.load_masters(&mut |_: u64| true);
    library.load_versions(&mut |_: u64| true);

    println!("Audit: {:?}", library.get_auditor().unwrap());
}

/// print the keywords with indentation for the hierarchy
fn print_keywords(keywords: &Vec<Keyword>, indent: &str) {
    for keyword in keywords {
        if !keyword.is_valid() {
            continue;
        }
        let name = keyword.name.as_ref().unwrap();
        let uuid = keyword.uuid().as_ref().unwrap();
        let parent = if let &Some(ref p) = keyword.parent() {
            p.clone()
        } else {
            "".to_owned()
        };
        println!("| {}{} | {} | {} |", indent, name,
                 uuid, parent);
        if keyword.children.is_some() {
            let new_indent;
            if indent.is_empty() {
                new_indent = "+- ".to_owned() + indent;
            } else {
                new_indent = "\t".to_owned() + indent;
            }
            print_keywords(keyword.children.as_ref().unwrap(), &new_indent);
        }
    }
}

fn process_dump(args: &Args) {
    let mut library = Library::new(&args.arg_path);

    {
        let version = library.library_version();
        println!("Version {}", version);
    }

    let model_info = library.get_model_info().unwrap();
    println!("model info");
    println!("\tDB version: {}", model_info.db_version.unwrap_or(0));
    println!("\tDB minor version: {}", model_info.db_minor_version.unwrap_or(0));
    println!("\tDB back compat: {}", model_info.db_minor_back_compatible_version.unwrap_or(0));
    println!("\tProject version: {}", model_info.project_version.unwrap_or(0));
    println!("\tCreation date: {}", model_info.create_date.unwrap_or("NONE".to_owned()));
    println!("\tImageIO: {} Camera RAW: {}",
             model_info.image_io_version.unwrap_or("NONE".to_owned()),
             model_info.raw_camera_bundle_version.unwrap_or("NONE".to_owned()));

    if args.flag_all || args.flag_folders {

        let mut pb = ProgressBar::new(1);
        pb.tick_format("|/-\\");

        library.load_folders(&mut |_: u64| {
            pb.tick();
            true
        });
        pb.finish();

        let folders = library.get_folders();
        println!("{} Folders:", folders.len());
        println!("| Name | uuid | impl album | type | model id | path |");
        for folder_uuid in folders {
            if folder_uuid.is_empty() {
                continue;
            }
            match library.get(folder_uuid) {
                Some(&StoreWrapper::Folder(ref folder)) => {
                    let name = folder.name.as_ref().unwrap();
                    let uuid = folder.uuid().as_ref().unwrap();
                    let implicit_album_uuid = if let Some(value) =
                        folder.implicit_album_uuid.as_ref() {
                            value
                        } else {
                            ""
                        };
                    let path = folder.path.as_ref().unwrap();
                    let folder_type = if let Some(ref ft) = folder.folder_type {
                        FolderType::to_int(ft)
                    } else {
                        0
                    };
                    println!("| {} | {} | {} | {} | {} | {} |",
                             name, uuid, implicit_album_uuid,
                             folder_type, folder.model_id(), path)
                },
                _ => println!("folder {} not found", folder_uuid)
            }
        }
    }
    if args.flag_all || args.flag_albums {

        let mut pb = ProgressBar::new(1);
        pb.tick_format("|/-\\");

        library.load_albums(&mut |_: u64| {
            pb.tick();
            true
        });
        pb.finish();

        let albums = library.get_albums();
        println!("{} Albums:", albums.len());
        println!("| name | uuid | parent (fldr) | query (fldr) | type | class | model id |");
        for album_uuid in albums {
            if album_uuid.is_empty() {
                continue;
            }
            match library.get(album_uuid) {
                Some(&StoreWrapper::Album(ref album)) => {
                    let name = if let Some(ref n) = album.name {
                        n.clone()
                    } else {
                        "".to_owned()
                    };
                    let uuid = album.uuid().as_ref().unwrap();
                    let parent = if let &Some(ref p) = album.parent() {
                        p.clone()
                    } else {
                        "".to_owned()
                    };
                    let query_folder_uuid = if let Some(ref qf) =
                        album.query_folder_uuid {
                            qf.clone()
                        } else {
                            "".to_owned()
                        };
                    let album_class = if let Some(ref c) = album.subclass {
                        AlbumSubclass::to_int(c)
                    } else {
                        0
                    };
                    println!("| {} | {} | {} | {} | {} | {} | {} |",
                             name, uuid, parent, query_folder_uuid,
                             album.album_type.unwrap_or(0),
                             album_class, album.model_id())
                },
                _ => println!("album {} not found", album_uuid)
            }
        }
    }
    if args.flag_all || args.flag_keywords {
        if let Some(ref keywords) = library.list_keywords() {
            println!("{} keywords:", keywords.len());
            println!("| name | uuid | parent |");
            print_keywords(keywords, &"".to_owned());
        }
    }

    if args.flag_all || args.flag_masters  {

        let count = model_info.master_count.unwrap_or(0) as u64;
        let mut pb = ProgressBar::new(count);

        library.load_masters(&mut |inc: u64| {
            pb.add(inc);
            true
        });
        pb.finish();

        let masters = library.get_masters();
        println!("{} Masters:", masters.len());
        println!("| uuid | project | path |");
        for master_uuid in masters {
            if master_uuid.is_empty() {
                continue;
            }
            match library.get(master_uuid) {
                Some(&StoreWrapper::Master(ref master)) => {
                    let uuid = master.uuid().as_ref().unwrap();
                    let parent = master.parent().as_ref().unwrap();
                    let image_path = master.image_path.as_ref().unwrap();
                    println!("| {} | {} | {} |", uuid, parent, image_path)
                },
                _ => println!("master {} not found", master_uuid)
            }
        }
    }
    if args.flag_all || args.flag_versions  {

        let count = model_info.version_count.unwrap_or(0) as u64;
        let mut pb = ProgressBar::new(count);

        library.load_versions(&mut |inc: u64| {
            pb.add(inc);
            true
        });
        pb.finish();

        let versions = library.get_versions();
        println!("{} Versions:", versions.len());
        println!("| uuid | master | project | name | original |");
        for version_uuid in versions {
            if version_uuid.is_empty() {
                continue;
            }
            match library.get(version_uuid) {
                Some(&StoreWrapper::Version(ref version)) => {
                    let uuid = version.uuid().as_ref().unwrap();
                    let parent = version.parent().as_ref().unwrap();
                    let project_uuid = version.project_uuid.as_ref().unwrap();
                    let name = version.name.as_ref().unwrap();

                    println!("| {} | {} | {} | {} | {} |",
                             uuid, parent, project_uuid, name,
                             version.is_original.unwrap_or(false))
                },
                _ => println!("version {} not found", version_uuid)
            }
        }

    }
}
<commit_msg>Improve output the dumper<commit_after>/*
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

extern crate aplib;
extern crate docopt;
extern crate rustc_serialize;
extern crate pbr;

use docopt::Docopt;
use rustc_serialize::{Decodable, Decoder};
use std::io::stderr;
use pbr::ProgressBar;

use aplib::AplibObject;
use aplib::Library;
use aplib::Keyword;
use aplib::StoreWrapper;
use aplib::audit::Reporter;
use aplib::AlbumSubclass;
use aplib::FolderType;

const USAGE: &'static str = "
Usage:
  dumper <command> ([--all] | [--albums] [--versions] [--masters] [--folders] [--keywords]) <path>

Options:
    --all          Select all objects
    --albums       Select only albums
    --masters      Select only masters
    --folders      Select only folders
    --keywords     Select only keywords

Commands are:
    dump           Dump the objects
    audit          Audit mode: output what we ignored
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_command: Command,
    arg_path: String,
    flag_all: bool,
    flag_albums: bool,
    flag_versions: bool,
    flag_masters: bool,
    flag_folders: bool,
    flag_keywords: bool
}

#[derive(Debug)]
enum Command {
    Dump,
    Audit,
    Unknown(String)
}

impl Decodable for Command {
    fn decode<D: Decoder>(d: &mut D) -> Result<Command, D::Error> {
        let s = try!(d.read_str());
        Ok(match &*s {
            "dump" => Command::Dump,
            "audit" => Command::Audit,
            s => Command::Unknown(s.to_owned()),
        })
    }
}

fn main() {

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    {
        match args.arg_command {
            Command::Dump =>
                process_dump(&args),
            Command::Audit =>
                process_audit(&args),
            _ =>
                ()
        };
    }
}

fn process_audit(args: &Args) {

    let mut library = Library::new(&args.arg_path);

    let auditor = Reporter::new();
    library.set_auditor(Some(auditor));

    library.library_version();
    library.load_folders(&mut |_: u64| true);
    library.load_albums(&mut |_: u64| true);
    library.load_masters(&mut |_: u64| true);
    library.load_versions(&mut |_: u64| true);

    println!("Audit: {:?}", library.get_auditor().unwrap());
}

/// print the keywords with indentation for the hierarchy
fn print_keywords(keywords: &Vec<Keyword>, indent: &str) {
    for keyword in keywords {
        if !keyword.is_valid() {
            continue;
        }
        let name = keyword.name.as_ref().unwrap();
        let uuid = keyword.uuid().as_ref().unwrap();
        let parent = if let &Some(ref p) = keyword.parent() {
            p.clone()
        } else {
            "".to_owned()
        };
        println!("| {:<26} | {:<26} | {}{}",
                 uuid, parent, indent, name);
        if keyword.children.is_some() {
            let new_indent;
            if indent.is_empty() {
                new_indent = "+- ".to_owned() + indent;
            } else {
                new_indent = "\t".to_owned() + indent;
            }
            print_keywords(keyword.children.as_ref().unwrap(), &new_indent);
        }
    }
}

fn process_dump(args: &Args) {
    let mut library = Library::new(&args.arg_path);

    {
        let version = library.library_version();
        println!("Version {}", version);
    }

    let model_info = library.get_model_info().unwrap();
    println!("model info");
    println!("\tDB version: {}", model_info.db_version.unwrap_or(0));
    println!("\tDB minor version: {}", model_info.db_minor_version.unwrap_or(0));
    println!("\tDB back compat: {}", model_info.db_minor_back_compatible_version.unwrap_or(0));
    println!("\tProject version: {}", model_info.project_version.unwrap_or(0));
    println!("\tCreation date: {}", model_info.create_date.unwrap_or("NONE".to_owned()));
    println!("\tImageIO: {} Camera RAW: {}",
             model_info.image_io_version.unwrap_or("NONE".to_owned()),
             model_info.raw_camera_bundle_version.unwrap_or("NONE".to_owned()));

    if args.flag_all || args.flag_folders {

        let mut pb = ProgressBar::on(stderr(), 1);
        pb.tick_format("|/-\\");

        library.load_folders(&mut |_: u64| {
            pb.tick();
            true
        });
        pb.finish();

        let folders = library.get_folders();
        println!("{} Folders:", folders.len());
        println!("| Name                       | uuid                       | impl album                 | type | model id | path");
        println!("+----------------------------+----------------------------+----------------------------+------+----------+----------");
        for folder_uuid in folders {
            if folder_uuid.is_empty() {
                continue;
            }
            match library.get(folder_uuid) {
                Some(&StoreWrapper::Folder(ref folder)) => {
                    let name = folder.name.as_ref().unwrap();
                    let uuid = folder.uuid().as_ref().unwrap();
                    let implicit_album_uuid = if let Some(value) =
                        folder.implicit_album_uuid.as_ref() {
                            value
                        } else {
                            ""
                        };
                    let path = folder.path.as_ref().unwrap();
                    let folder_type = if let Some(ref ft) = folder.folder_type {
                        FolderType::to_int(ft)
                    } else {
                        0
                    };
                    println!("| {:<26} | {:<26} | {:<26} | {:>4} | {:>8} | {}",
                             name, uuid, implicit_album_uuid,
                             folder_type, folder.model_id(), path)
                },
                _ => println!("folder {} not found", folder_uuid)
            }
        }
    }
    if args.flag_all || args.flag_albums {

        let mut pb = ProgressBar::on(stderr(), 1);
        pb.tick_format("|/-\\");

        library.load_albums(&mut |_: u64| {
            pb.tick();
            true
        });
        pb.finish();

        let albums = library.get_albums();
        println!("{} Albums:", albums.len());
        println!("| uuid                       | parent (fldr)              | query (fldr)               | type | class | model id | name");
        println!("+----------------------------+----------------------------+----------------------------+------+-------+----------+-----");
        for album_uuid in albums {
            if album_uuid.is_empty() {
                continue;
            }
            match library.get(album_uuid) {
                Some(&StoreWrapper::Album(ref album)) => {
                    let name = if let Some(ref n) = album.name {
                        n.clone()
                    } else {
                        "".to_owned()
                    };
                    let uuid = album.uuid().as_ref().unwrap();
                    let parent = if let &Some(ref p) = album.parent() {
                        p.clone()
                    } else {
                        "".to_owned()
                    };
                    let query_folder_uuid = if let Some(ref qf) =
                        album.query_folder_uuid {
                            qf.clone()
                        } else {
                            "".to_owned()
                        };
                    let album_class = if let Some(ref c) = album.subclass {
                        AlbumSubclass::to_int(c)
                    } else {
                        0
                    };
                    println!("| {:<26} | {:<26} | {:<26} | {:>4} | {:>5} | {:>8} | {}",
                             uuid, parent, query_folder_uuid,
                             album.album_type.unwrap_or(0),
                             album_class, album.model_id(), name)
                },
                _ => println!("album {} not found", album_uuid)
            }
        }
    }
    if args.flag_all || args.flag_keywords {
        if let Some(ref keywords) = library.list_keywords() {
            println!("{} keywords:", keywords.len());
            println!("| uuid                       | parent                     | name");
            println!("+----------------------------+----------------------------+-----------");
            print_keywords(keywords, &"".to_owned());
        }
    }

    if args.flag_all || args.flag_masters  {

        let count = model_info.master_count.unwrap_or(0) as u64;
        let mut pb = ProgressBar::on(stderr(), count);

        library.load_masters(&mut |inc: u64| {
            pb.add(inc);
            true
        });
        pb.finish();

        let masters = library.get_masters();
        println!("{} Masters:", masters.len());
        println!("| uuid                       | project                    | path");
        println!("+----------------------------+----------------------------+-----------------------");
        for master_uuid in masters {
            if master_uuid.is_empty() {
                continue;
            }
            match library.get(master_uuid) {
                Some(&StoreWrapper::Master(ref master)) => {
                    let uuid = master.uuid().as_ref().unwrap();
                    let parent = master.parent().as_ref().unwrap();
                    let image_path = master.image_path.as_ref().unwrap();
                    println!("| {:<26} | {:<26} | {}", uuid, parent, image_path)
                },
                _ => println!("master {} not found", master_uuid)
            }
        }
    }
    if args.flag_all || args.flag_versions  {

        let count = model_info.version_count.unwrap_or(0) as u64;
        let mut pb = ProgressBar::on(stderr(), count);

        library.load_versions(&mut |inc: u64| {
            pb.add(inc);
            true
        });
        pb.finish();

        let versions = library.get_versions();
        println!("{} Versions:", versions.len());
        println!("| uuid                       | master                     | project                    | original | name");
        println!("+----------------------------+----------------------------+----------------------------+----------+------------");
        for version_uuid in versions {
            if version_uuid.is_empty() {
                continue;
            }
            match library.get(version_uuid) {
                Some(&StoreWrapper::Version(ref version)) => {
                    let uuid = version.uuid().as_ref().unwrap();
                    let parent = version.parent().as_ref().unwrap();
                    let project_uuid = version.project_uuid.as_ref().unwrap();
                    let name = version.name.as_ref().unwrap();

                    println!("| {:<26} | {:<26} | {:<26} | {:>8} | {}",
                             uuid, parent, project_uuid,
                             version.is_original.unwrap_or(false),
                             name)
                },
                _ => println!("version {} not found", version_uuid)
            }
        }

    }
}
<|endoftext|>


<commit_before>extern mod opengles;   // FIXME: Should only be for tests.
use glut::{check_loop, create_window, destroy_window, init, init_display_mode, swap_buffers, GLint};
use self::opengles::gl2::{ARRAY_BUFFER, COLOR_BUFFER_BIT, COMPILE_STATUS};
use self::opengles::gl2::{FRAGMENT_SHADER, LINK_STATUS, NO_ERROR, STATIC_DRAW};
use self::opengles::gl2::{TRIANGLE_STRIP, VERTEX_SHADER, GLclampf, GLenum};
use self::opengles::gl2::{GLsizei, GLuint, attach_shader, bind_buffer};
use self::opengles::gl2::{buffer_data, create_program, clear, clear_color};
use self::opengles::gl2::{compile_shader, create_shader, draw_arrays};
use self::opengles::gl2::{enable_vertex_attrib_array, gen_buffers};
use self::opengles::gl2::{get_attrib_location, get_error, get_program_iv};
use self::opengles::gl2::{get_shader_info_log, get_shader_iv};
use self::opengles::gl2::{get_uniform_location, link_program, shader_source};
use self::opengles::gl2::{use_program, vertex_attrib_pointer_f32};

use libc::{c_int, c_uint};
use oldcomm::{Chan, peek, Port, recv, send};
use io::println;
use ptr::{addr_of, null};
use str::to_bytes;
use task::TaskBuilder;
use vec::raw::to_ptr;

fn fragment_shader_source() -> ~str {
    ~"
    #ifdef GLES2
        precision mediump float;
    #endif

        void main(void) {
            gl_FragColor = vec4(1.0, 1.0, 1.0, 1.0);
        }
    "
}

fn vertex_shader_source() -> ~str {
    ~"
        attribute vec3 aVertexPosition;

        /*uniform mat4 uMVMatrix;
        uniform mat4 uPMatrix;*/

        void main(void) {
            gl_Position = /*uPMatrix * uMVMatrix **/
                vec4(aVertexPosition, 1.0);
        }
    "
}

fn load_shader(source_str: ~str, shader_type: GLenum) -> GLuint {
    let shader_id = create_shader(shader_type);
    shader_source(shader_id, ~[to_bytes(source_str)]);
    compile_shader(shader_id);

    if get_error() != NO_ERROR {
        println(fmt!("error: %d", get_error() as int));
        fail!(~"failed to compile shader with error");
    }

    if get_shader_iv(shader_id, COMPILE_STATUS) == (0 as GLint) {
        println(get_shader_info_log(shader_id));
        fail!(~"failed to compile shader");
    }
    return shader_id;
}

struct ShaderProgram {
    program: GLuint,
    aVertexPosition: c_int,
    /*let uPMatrix: c_int;
    let uMVMatrix: c_int;*/
}

fn ShaderProgram(program: GLuint) -> ShaderProgram {
    let p = ShaderProgram {
        program : program,
        aVertexPosition : get_attrib_location(program, ~"aVertexPosition"),
        /*self.uPMatrix : get_uniform_location(program, "uPMatrix"),
        self.uMVMatrix : get_uniform_location(program, "uMVMatrix")*/
    };
    enable_vertex_attrib_array(p.aVertexPosition as GLuint);
    return p;
}

fn init_shaders() -> ShaderProgram {
    let vertex_shader = load_shader(vertex_shader_source(), VERTEX_SHADER);
    let fragment_shader = load_shader(fragment_shader_source(),
                                      FRAGMENT_SHADER);

    let program = create_program();
    attach_shader(program, vertex_shader);
    attach_shader(program, fragment_shader);
    link_program(program);

    if get_program_iv(program, LINK_STATUS) == (0 as GLint) {
        fail!(~"failed to initialize program");
    }

    use_program(program);

    return ShaderProgram(program);
}

fn init_buffers() -> GLuint {
    let triangle_vertex_buffer = gen_buffers(1 as GLsizei)[0];
    bind_buffer(ARRAY_BUFFER, triangle_vertex_buffer);
    let vertices = ~[
        0.0f32, 1.0f32, 0.0f32,
        1.0f32, 0.0f32, 0.0f32,
        0.0f32, 0.0f32, 0.0f32
    ];
    buffer_data(ARRAY_BUFFER, vertices, STATIC_DRAW);
    return triangle_vertex_buffer;
}

fn draw_scene(shader_program: ShaderProgram, vertex_buffer: GLuint) {
    clear_color(0.0f32, 0.0f32, 1.0f32, 1.0f32);
    clear(COLOR_BUFFER_BIT);

    bind_buffer(ARRAY_BUFFER, vertex_buffer);
    vertex_attrib_pointer_f32(shader_program.aVertexPosition as GLuint,
                              3 as GLint, false, 0 as GLsizei, 0 as GLuint);
    draw_arrays(TRIANGLE_STRIP, 0 as GLint, 3 as GLint);
}

fn display_callback() {
    let program = init_shaders();
    let vertex_buffer = init_buffers();
    draw_scene(program, vertex_buffer);

    swap_buffers();
}

/*
#[test]
fn test_triangle_and_square() unsafe {
    let builder = task::task().sched_mode(task::PlatformThread);

    let po: Port<()> = Port();
    let ch = Chan(&po);
    let _result_ch: Chan<()> = builder.spawn_listener(|_port| {
        init();
        init_display_mode(0 as c_uint);
        let window = create_window(~"Rust GLUT");
        display_func(display_callback);

        let wakeup: Port<()> = Port();
        let wakeup_chan = Chan(&wakeup);
        timer_func(1000, || send(wakeup_chan, ()));

        loop {
            check_loop();

            if peek(wakeup) {
                recv(wakeup);
                send(ch, ());
                destroy_window(window);
                break;
            }
        }
    });

    recv(po);
}

*/
<commit_msg>Update tests for language changes<commit_after>extern mod opengles;   // FIXME: Should only be for tests.
use glut::{check_loop, create_window, destroy_window, init, init_display_mode, swap_buffers, GLint};
use self::opengles::gl2::{ARRAY_BUFFER, COLOR_BUFFER_BIT, COMPILE_STATUS};
use self::opengles::gl2::{FRAGMENT_SHADER, LINK_STATUS, NO_ERROR, STATIC_DRAW};
use self::opengles::gl2::{TRIANGLE_STRIP, VERTEX_SHADER, GLclampf, GLenum};
use self::opengles::gl2::{GLsizei, GLuint, attach_shader, bind_buffer};
use self::opengles::gl2::{buffer_data, create_program, clear, clear_color};
use self::opengles::gl2::{compile_shader, create_shader, draw_arrays};
use self::opengles::gl2::{enable_vertex_attrib_array, gen_buffers};
use self::opengles::gl2::{get_attrib_location, get_error, get_program_iv};
use self::opengles::gl2::{get_shader_info_log, get_shader_iv};
use self::opengles::gl2::{get_uniform_location, link_program, shader_source};
use self::opengles::gl2::{use_program, vertex_attrib_pointer_f32};

use core::pipes::{Chan, Port};
use core::pipes;
use libc::{c_int, c_uint};
use io::println;
use ptr::{addr_of, null};
use str::to_bytes;
use task::TaskBuilder;
use vec::raw::to_ptr;

fn fragment_shader_source() -> ~str {
    ~"
    #ifdef GLES2
        precision mediump float;
    #endif

        void main(void) {
            gl_FragColor = vec4(1.0, 1.0, 1.0, 1.0);
        }
    "
}

fn vertex_shader_source() -> ~str {
    ~"
        attribute vec3 aVertexPosition;

        /*uniform mat4 uMVMatrix;
        uniform mat4 uPMatrix;*/

        void main(void) {
            gl_Position = /*uPMatrix * uMVMatrix **/
                vec4(aVertexPosition, 1.0);
        }
    "
}

fn load_shader(source_str: ~str, shader_type: GLenum) -> GLuint {
    let shader_id = create_shader(shader_type);
    shader_source(shader_id, ~[to_bytes(source_str)]);
    compile_shader(shader_id);

    if get_error() != NO_ERROR {
        println(fmt!("error: %d", get_error() as int));
        fail!(~"failed to compile shader with error");
    }

    if get_shader_iv(shader_id, COMPILE_STATUS) == (0 as GLint) {
        println(get_shader_info_log(shader_id));
        fail!(~"failed to compile shader");
    }
    return shader_id;
}

struct ShaderProgram {
    program: GLuint,
    aVertexPosition: c_int,
    /*let uPMatrix: c_int;
    let uMVMatrix: c_int;*/
}

fn ShaderProgram(program: GLuint) -> ShaderProgram {
    let p = ShaderProgram {
        program : program,
        aVertexPosition : get_attrib_location(program, ~"aVertexPosition"),
        /*self.uPMatrix : get_uniform_location(program, "uPMatrix"),
        self.uMVMatrix : get_uniform_location(program, "uMVMatrix")*/
    };
    enable_vertex_attrib_array(p.aVertexPosition as GLuint);
    return p;
}

fn init_shaders() -> ShaderProgram {
    let vertex_shader = load_shader(vertex_shader_source(), VERTEX_SHADER);
    let fragment_shader = load_shader(fragment_shader_source(),
                                      FRAGMENT_SHADER);

    let program = create_program();
    attach_shader(program, vertex_shader);
    attach_shader(program, fragment_shader);
    link_program(program);

    if get_program_iv(program, LINK_STATUS) == (0 as GLint) {
        fail!(~"failed to initialize program");
    }

    use_program(program);

    return ShaderProgram(program);
}

fn init_buffers() -> GLuint {
    let triangle_vertex_buffer = gen_buffers(1 as GLsizei)[0];
    bind_buffer(ARRAY_BUFFER, triangle_vertex_buffer);
    let vertices = ~[
        0.0f32, 1.0f32, 0.0f32,
        1.0f32, 0.0f32, 0.0f32,
        0.0f32, 0.0f32, 0.0f32
    ];
    buffer_data(ARRAY_BUFFER, vertices, STATIC_DRAW);
    return triangle_vertex_buffer;
}

fn draw_scene(shader_program: ShaderProgram, vertex_buffer: GLuint) {
    clear_color(0.0f32, 0.0f32, 1.0f32, 1.0f32);
    clear(COLOR_BUFFER_BIT);

    bind_buffer(ARRAY_BUFFER, vertex_buffer);
    vertex_attrib_pointer_f32(shader_program.aVertexPosition as GLuint,
                              3 as GLint, false, 0 as GLsizei, 0 as GLuint);
    draw_arrays(TRIANGLE_STRIP, 0 as GLint, 3 as GLint);
}

fn display_callback() {
    let program = init_shaders();
    let vertex_buffer = init_buffers();
    draw_scene(program, vertex_buffer);

    swap_buffers();
}

/*
#[test]
fn test_triangle_and_square() unsafe {
    let builder = task::task().sched_mode(task::PlatformThread);

    let po: Port<()> = Port();
    let ch = Chan(&po);
    let _result_ch: Chan<()> = builder.spawn_listener(|_port| {
        init();
        init_display_mode(0 as c_uint);
        let window = create_window(~"Rust GLUT");
        display_func(display_callback);

        let wakeup: Port<()> = Port();
        let wakeup_chan = Chan(&wakeup);
        timer_func(1000, || send(wakeup_chan, ()));

        loop {
            check_loop();

            if peek(wakeup) {
                recv(wakeup);
                send(ch, ());
                destroy_window(window);
                break;
            }
        }
    });

    recv(po);
}

*/
<|endoftext|>


<commit_before>// Copyright 2013-2014 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_id="cgmath#0.1"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![comment = "A mathematics library for computer graphics."]
#![license = "ASL2"]

#![feature(globs)]
#![feature(macro_rules)]

//! Computer graphics-centric math.
//!
//! This crate provides useful mathematical primitives and operations on them.
//! It is organized into one module per primitive. The core structures are
//! vectors and matrices. A strongly-typed interface is provided, to prevent
//! mixing units or violating mathematical invariants.
//!
//! Transformations are not usually done directly on matrices, but go through
//! transformation objects that can be converted to matrices. Rotations go
//! through the `Basis` types, which are guaranteed to be orthogonal matrices.
//! Despite this, one can directly create a limited rotation matrix using the
//! `look_at`, `from_angle`, `from_euler`, and `from_axis_angle` methods.
//! These are provided for convenience.

pub mod array;

pub mod matrix;
pub mod quaternion;
pub mod vector;

pub mod angle;
pub mod plane;
pub mod point;
pub mod line;
pub mod ray;
pub mod rotation;
pub mod transform;

pub mod projection;

pub mod aabb;
pub mod cylinder;
pub mod frustum;
pub mod intersect;
pub mod obb;
pub mod sphere;

pub mod approx;
pub mod num;
<commit_msg>Updated to latest Rust: Added 'extern crate debug;'<commit_after>// Copyright 2013-2014 The CGMath Developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![crate_id="cgmath#0.1"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![comment = "A mathematics library for computer graphics."]
#![license = "ASL2"]

#![feature(globs)]
#![feature(macro_rules)]

//! Computer graphics-centric math.
//!
//! This crate provides useful mathematical primitives and operations on them.
//! It is organized into one module per primitive. The core structures are
//! vectors and matrices. A strongly-typed interface is provided, to prevent
//! mixing units or violating mathematical invariants.
//!
//! Transformations are not usually done directly on matrices, but go through
//! transformation objects that can be converted to matrices. Rotations go
//! through the `Basis` types, which are guaranteed to be orthogonal matrices.
//! Despite this, one can directly create a limited rotation matrix using the
//! `look_at`, `from_angle`, `from_euler`, and `from_axis_angle` methods.
//! These are provided for convenience.

extern crate debug;

pub mod array;

pub mod matrix;
pub mod quaternion;
pub mod vector;

pub mod angle;
pub mod plane;
pub mod point;
pub mod line;
pub mod ray;
pub mod rotation;
pub mod transform;

pub mod projection;

pub mod aabb;
pub mod cylinder;
pub mod frustum;
pub mod intersect;
pub mod obb;
pub mod sphere;

pub mod approx;
pub mod num;
<|endoftext|>


<commit_before>use common::{CommonResources, CommonOperatorOutput};
use kernels::softmax::{SoftmaxKernel};

use float::ord::{F32InfNan};
use iter_utils::{argmax}; //, KahanSum};
use operator::prelude::*;
use rng::xorshift::{Xorshiftplus128Rng};

use std::u32;

#[derive(Clone, Copy)]
pub struct ClassLossConfig {
  pub batch_sz:     usize,
  pub num_classes:  usize,
}

pub struct SoftmaxNLLClassLossOperator {
  cfg:      ClassLossConfig,
  in_:      CommonOperatorOutput<f32>,
  /*max_log:  Vec<f32>,
  facts:    Vec<f32>,
  sum_fact: Vec<f32>,*/
  hats:     Vec<u32>,
  losses:   Vec<f32>,
  loss1:    f32,
  accuracy: usize,
  labels:   Vec<u32>,
  weights:  Vec<f32>,
  sm_kern:  SoftmaxKernel,
  out:      CommonOperatorOutput<f32>,
}

impl SoftmaxNLLClassLossOperator {
  pub fn new(cfg: ClassLossConfig, cap: OpCapability, prev_op: &DiffOperator<f32, Output=CommonOperatorOutput<f32>, Rng=Xorshiftplus128Rng>, prev_arm: usize, res: CommonResources) -> SoftmaxNLLClassLossOperator {
    /*let mut max_log = Vec::with_capacity(cfg.batch_sz);
    unsafe { max_log.set_len(cfg.batch_sz) };
    let mut facts = Vec::with_capacity(cfg.batch_sz * cfg.num_classes);
    unsafe { facts.set_len(cfg.batch_sz * cfg.num_classes) };
    let mut sum_fact = Vec::with_capacity(cfg.batch_sz);
    unsafe { sum_fact.set_len(cfg.batch_sz) };*/
    let mut hats = Vec::with_capacity(cfg.batch_sz);
    hats.resize(cfg.batch_sz, 0);
    let mut losses = Vec::with_capacity(cfg.batch_sz);
    losses.resize(cfg.batch_sz, 0.0);
    let mut labels = Vec::with_capacity(cfg.batch_sz);
    labels.resize(cfg.batch_sz, 0);
    let mut weights = Vec::with_capacity(cfg.batch_sz);
    weights.resize(cfg.batch_sz, 0.0);
    SoftmaxNLLClassLossOperator{
      cfg:      cfg,
      in_:      prev_op._output(prev_arm),
      /*max_log:  max_log,
      facts:    facts,
      sum_fact: sum_fact,*/
      hats:     hats,
      losses:   losses,
      loss1:    0.0,
      accuracy: 0,
      labels:   labels,
      weights:  weights,
      sm_kern:  SoftmaxKernel::new(cfg.batch_sz, cfg.num_classes, res.nnp_pool),
      out:      CommonOperatorOutput::new(cfg.batch_sz, cfg.num_classes, cap),
    }
  }
}

impl<S> DiffOperatorInput<f32, S> for SoftmaxNLLClassLossOperator where S: SampleLabel + SampleLossWeight<ClassLoss> {
  fn load_data(&mut self, samples: &[S]) {
    let actual_batch_size = samples.len();
    assert!(actual_batch_size <= self.cfg.batch_sz);
    for (idx, sample) in samples.iter().enumerate() {
      if let Some(cat) = sample.class() {
        assert!(cat < self.cfg.num_classes as u32);
        self.labels[idx] = cat;
      } else {
        self.labels[idx] = u32::MAX;
      }
      self.weights[idx] = sample.weight().unwrap_or(1.0);
    }
    *self.out.batch_size.borrow_mut() = actual_batch_size;
  }
}

impl DiffOperator<f32> for SoftmaxNLLClassLossOperator {
  type Output = CommonOperatorOutput<f32>;
  type Rng = Xorshiftplus128Rng;

  fn _output(&self, _arm: usize) -> CommonOperatorOutput<f32> {
    assert_eq!(0, _arm);
    self.out.clone()
  }

  fn reset_loss(&mut self) {
    self.loss1 = 0.0;
    self.accuracy = 0;
  }

  fn store_loss(&mut self) -> f32 {
    self.loss1
  }

  fn forward(&mut self, _phase: OpPhase) {
    //self.out.batch_size = self.in_.batch_size;
    let batch_size = *self.in_.batch_size.borrow();
    assert_eq!(batch_size, *self.out.batch_size.borrow());

    /*for idx in 0 .. batch_size {
      let range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(self.in_.out_buf.borrow()[range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      let max_logit = self.in_.out_buf.borrow()[idx * self.cfg.num_classes + max_logit_k];
      self.max_log[idx] = max_logit;
      self.hats[idx] = max_logit_k as u32;
      for k in 0 .. self.cfg.num_classes {
        self.facts[idx * self.cfg.num_classes + k] = (self.in_.out_buf.borrow()[idx * self.cfg.num_classes + k] - max_logit).exp();
      }
      let sum_fact: f32 = Sum::sum(self.facts[range].iter().map(|&x| x));
      for k in 0 .. self.cfg.num_classes {
        self.out.out_buf.borrow_mut()[idx * self.cfg.num_classes + k] = self.facts[idx * self.cfg.num_classes + k] / sum_fact;
      }
      self.losses[idx] = -self.weights[idx] * self.out.out_buf.borrow()[idx * self.cfg.num_classes + self.labels[idx] as usize].ln();
    }*/
    let in_buf = self.in_.out_buf.borrow();
    //println!("DEBUG: softmax: in buf: {:?}", &in_buf[ .. 10]);
    let mut out_buf = self.out.out_buf.borrow_mut();
    self.sm_kern.forward(batch_size, &*in_buf, &mut *out_buf);
    let mut batch_loss = 0.0;
    let mut batch_accuracy = 0;
    for idx in 0 .. batch_size {
      let idx_range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(in_buf[idx_range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      self.hats[idx] = max_logit_k as u32;
      if self.hats[idx] == self.labels[idx] {
        batch_accuracy += 1;
      }
      let loss = -self.weights[idx] * out_buf[idx * self.cfg.num_classes + self.labels[idx] as usize].ln();
      self.losses[idx] = loss;
      batch_loss += loss;
    }
    //self.loss1 += Sum::sum(self.losses[ .. batch_size].iter().map(|&x| x));
    //println!("DEBUG: softmax: out buf: {:?}", &out_buf[ .. 10]);
    //println!("DEBUG: softmax: accuracy: {}/{}", batch_accuracy, batch_size);

    let in_loss = *self.in_.out_loss.borrow();
    //self.loss1 += batch_loss + in_loss;
    self.loss1 += batch_loss;
    self.accuracy += batch_accuracy;
    *self.out.out_loss.borrow_mut() = batch_loss + in_loss;
  }

  fn backward(&mut self) {
    //assert_eq!(self.out.batch_size, self.in_.batch_size);
    let batch_size = *self.out.batch_size.borrow();
    if let Some(ref mut in_grad) = self.in_.out_grad.as_mut() {
      let out_buf = self.out.out_buf.borrow();
      let mut in_grad = in_grad.borrow_mut();
      let mut p = 0;
      for idx in 0 .. batch_size {
        for k in 0 .. self.cfg.num_classes {
          in_grad[p] =
              self.weights[idx] *
              (out_buf[p] - if k == self.labels[idx] as usize { 1.0 } else { 0.0 });
          p += 1;
        }
      }
    }
  }
}
<commit_msg>Implement entropy-regularized softmax NLL loss.<commit_after>use common::{CommonResources, CommonOperatorOutput};
use kernels::softmax::{SoftmaxKernel};

use float::ord::{F32InfNan};
use iter_utils::{argmax}; //, KahanSum};
use operator::prelude::*;
use rng::xorshift::{Xorshiftplus128Rng};

use std::u32;

#[derive(Clone, Copy)]
pub struct ClassLossConfig {
  pub batch_sz:     usize,
  pub num_classes:  usize,
}

pub struct SoftmaxNLLClassLossOperator {
  cfg:      ClassLossConfig,
  in_:      CommonOperatorOutput<f32>,
  /*max_log:  Vec<f32>,
  facts:    Vec<f32>,
  sum_fact: Vec<f32>,*/
  hats:     Vec<u32>,
  losses:   Vec<f32>,
  loss1:    f32,
  accuracy: usize,
  labels:   Vec<u32>,
  weights:  Vec<f32>,
  sm_kern:  SoftmaxKernel,
  out:      CommonOperatorOutput<f32>,
}

impl SoftmaxNLLClassLossOperator {
  pub fn new(cfg: ClassLossConfig, cap: OpCapability, prev_op: &DiffOperator<f32, Output=CommonOperatorOutput<f32>, Rng=Xorshiftplus128Rng>, prev_arm: usize, res: CommonResources) -> SoftmaxNLLClassLossOperator {
    /*let mut max_log = Vec::with_capacity(cfg.batch_sz);
    unsafe { max_log.set_len(cfg.batch_sz) };
    let mut facts = Vec::with_capacity(cfg.batch_sz * cfg.num_classes);
    unsafe { facts.set_len(cfg.batch_sz * cfg.num_classes) };
    let mut sum_fact = Vec::with_capacity(cfg.batch_sz);
    unsafe { sum_fact.set_len(cfg.batch_sz) };*/
    let mut hats = Vec::with_capacity(cfg.batch_sz);
    hats.resize(cfg.batch_sz, 0);
    let mut losses = Vec::with_capacity(cfg.batch_sz);
    losses.resize(cfg.batch_sz, 0.0);
    let mut labels = Vec::with_capacity(cfg.batch_sz);
    labels.resize(cfg.batch_sz, 0);
    let mut weights = Vec::with_capacity(cfg.batch_sz);
    weights.resize(cfg.batch_sz, 0.0);
    SoftmaxNLLClassLossOperator{
      cfg:      cfg,
      in_:      prev_op._output(prev_arm),
      /*max_log:  max_log,
      facts:    facts,
      sum_fact: sum_fact,*/
      hats:     hats,
      losses:   losses,
      loss1:    0.0,
      accuracy: 0,
      labels:   labels,
      weights:  weights,
      sm_kern:  SoftmaxKernel::new(cfg.batch_sz, cfg.num_classes, res.nnp_pool),
      out:      CommonOperatorOutput::new(cfg.batch_sz, cfg.num_classes, cap),
    }
  }
}

impl<S> DiffOperatorInput<f32, S> for SoftmaxNLLClassLossOperator where S: SampleLabel + SampleLossWeight<ClassLoss> {
  fn load_data(&mut self, samples: &[S]) {
    let actual_batch_size = samples.len();
    assert!(actual_batch_size <= self.cfg.batch_sz);
    for (idx, sample) in samples.iter().enumerate() {
      if let Some(cat) = sample.class() {
        assert!(cat < self.cfg.num_classes as u32);
        assert!(cat != u32::MAX);
        self.labels[idx] = cat;
      } else {
        self.labels[idx] = u32::MAX;
      }
      self.weights[idx] = sample.weight().unwrap_or(1.0);
    }
    *self.out.batch_size.borrow_mut() = actual_batch_size;
  }
}

impl DiffOperator<f32> for SoftmaxNLLClassLossOperator {
  type Output = CommonOperatorOutput<f32>;
  type Rng = Xorshiftplus128Rng;

  fn _output(&self, _arm: usize) -> CommonOperatorOutput<f32> {
    assert_eq!(0, _arm);
    self.out.clone()
  }

  fn reset_loss(&mut self) {
    self.loss1 = 0.0;
    self.accuracy = 0;
  }

  fn store_loss(&mut self) -> f32 {
    self.loss1
  }

  fn forward(&mut self, _phase: OpPhase) {
    //self.out.batch_size = self.in_.batch_size;
    let batch_size = *self.in_.batch_size.borrow();
    assert_eq!(batch_size, *self.out.batch_size.borrow());

    /*for idx in 0 .. batch_size {
      let range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(self.in_.out_buf.borrow()[range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      let max_logit = self.in_.out_buf.borrow()[idx * self.cfg.num_classes + max_logit_k];
      self.max_log[idx] = max_logit;
      self.hats[idx] = max_logit_k as u32;
      for k in 0 .. self.cfg.num_classes {
        self.facts[idx * self.cfg.num_classes + k] = (self.in_.out_buf.borrow()[idx * self.cfg.num_classes + k] - max_logit).exp();
      }
      let sum_fact: f32 = Sum::sum(self.facts[range].iter().map(|&x| x));
      for k in 0 .. self.cfg.num_classes {
        self.out.out_buf.borrow_mut()[idx * self.cfg.num_classes + k] = self.facts[idx * self.cfg.num_classes + k] / sum_fact;
      }
      self.losses[idx] = -self.weights[idx] * self.out.out_buf.borrow()[idx * self.cfg.num_classes + self.labels[idx] as usize].ln();
    }*/
    let in_buf = self.in_.out_buf.borrow();
    //println!("DEBUG: softmax: in buf: {:?}", &in_buf[ .. 10]);
    let mut out_buf = self.out.out_buf.borrow_mut();
    self.sm_kern.forward(batch_size, &*in_buf, &mut *out_buf);
    let mut batch_loss = 0.0;
    let mut batch_accuracy = 0;
    for idx in 0 .. batch_size {
      let idx_range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(in_buf[idx_range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      self.hats[idx] = max_logit_k as u32;
      if self.hats[idx] == self.labels[idx] {
        batch_accuracy += 1;
      }
      let loss = if self.labels[idx] == u32::MAX {
        0.0
      } else {
        -self.weights[idx] * out_buf[idx * self.cfg.num_classes + self.labels[idx] as usize].ln()
      };
      self.losses[idx] = loss;
      batch_loss += loss;
    }
    //self.loss1 += Sum::sum(self.losses[ .. batch_size].iter().map(|&x| x));
    //println!("DEBUG: softmax: out buf: {:?}", &out_buf[ .. 10]);
    //println!("DEBUG: softmax: accuracy: {}/{}", batch_accuracy, batch_size);

    let in_loss = *self.in_.out_loss.borrow();
    //self.loss1 += batch_loss + in_loss;
    self.loss1 += batch_loss;
    self.accuracy += batch_accuracy;
    *self.out.out_loss.borrow_mut() = batch_loss + in_loss;
  }

  fn backward(&mut self) {
    //assert_eq!(self.out.batch_size, self.in_.batch_size);
    let batch_size = *self.out.batch_size.borrow();
    if let Some(ref mut in_grad) = self.in_.out_grad.as_mut() {
      let out_buf = self.out.out_buf.borrow();
      let mut in_grad = in_grad.borrow_mut();
      let mut p = 0;
      for idx in 0 .. batch_size {
        for k in 0 .. self.cfg.num_classes {
          in_grad[p] =
              self.weights[idx] *
              (out_buf[p] - if k == self.labels[idx] as usize { 1.0 } else { 0.0 });
          p += 1;
        }
      }
    }
  }
}

#[derive(Clone, Copy)]
pub struct EntRegSoftmaxNLLClassLossConfig {
  pub batch_sz:     usize,
  pub num_classes:  usize,
  pub entropy_coef: f32,
}

pub struct EntRegSoftmaxNLLClassLossOperator {
  cfg:      EntRegSoftmaxNLLClassLossConfig,
  in_:      CommonOperatorOutput<f32>,
  /*max_log:  Vec<f32>,
  facts:    Vec<f32>,
  sum_fact: Vec<f32>,*/
  hats:     Vec<u32>,
  losses:   Vec<f32>,
  ents:     Vec<f32>,
  loss1:    f32,
  accuracy: usize,
  labels:   Vec<u32>,
  weights:  Vec<f32>,
  sm_kern:  SoftmaxKernel,
  out:      CommonOperatorOutput<f32>,
}

impl EntRegSoftmaxNLLClassLossOperator {
  pub fn new(cfg: EntRegSoftmaxNLLClassLossConfig, cap: OpCapability, prev_op: &DiffOperator<f32, Output=CommonOperatorOutput<f32>, Rng=Xorshiftplus128Rng>, prev_arm: usize, res: CommonResources) -> EntRegSoftmaxNLLClassLossOperator {
    /*let mut max_log = Vec::with_capacity(cfg.batch_sz);
    unsafe { max_log.set_len(cfg.batch_sz) };
    let mut facts = Vec::with_capacity(cfg.batch_sz * cfg.num_classes);
    unsafe { facts.set_len(cfg.batch_sz * cfg.num_classes) };
    let mut sum_fact = Vec::with_capacity(cfg.batch_sz);
    unsafe { sum_fact.set_len(cfg.batch_sz) };*/
    let mut hats = Vec::with_capacity(cfg.batch_sz);
    hats.resize(cfg.batch_sz, 0);
    let mut losses = Vec::with_capacity(cfg.batch_sz);
    losses.resize(cfg.batch_sz, 0.0);
    let mut ents = Vec::with_capacity(cfg.batch_sz);
    ents.resize(cfg.batch_sz, 0.0);
    let mut labels = Vec::with_capacity(cfg.batch_sz);
    labels.resize(cfg.batch_sz, 0);
    let mut weights = Vec::with_capacity(cfg.batch_sz);
    weights.resize(cfg.batch_sz, 0.0);
    EntRegSoftmaxNLLClassLossOperator{
      cfg:      cfg,
      in_:      prev_op._output(prev_arm),
      /*max_log:  max_log,
      facts:    facts,
      sum_fact: sum_fact,*/
      hats:     hats,
      losses:   losses,
      ents:     ents,
      loss1:    0.0,
      accuracy: 0,
      labels:   labels,
      weights:  weights,
      sm_kern:  SoftmaxKernel::new(cfg.batch_sz, cfg.num_classes, res.nnp_pool),
      out:      CommonOperatorOutput::new(cfg.batch_sz, cfg.num_classes, cap),
    }
  }
}

impl<S> DiffOperatorInput<f32, S> for EntRegSoftmaxNLLClassLossOperator where S: SampleLabel + SampleLossWeight<ClassLoss> {
  fn load_data(&mut self, samples: &[S]) {
    let actual_batch_size = samples.len();
    assert!(actual_batch_size <= self.cfg.batch_sz);
    for (idx, sample) in samples.iter().enumerate() {
      if let Some(cat) = sample.class() {
        assert!(cat < self.cfg.num_classes as u32);
        assert!(cat != u32::MAX);
        self.labels[idx] = cat;
      } else {
        self.labels[idx] = u32::MAX;
      }
      self.weights[idx] = sample.weight().unwrap_or(1.0);
    }
    *self.out.batch_size.borrow_mut() = actual_batch_size;
  }
}

impl DiffOperator<f32> for EntRegSoftmaxNLLClassLossOperator {
  type Output = CommonOperatorOutput<f32>;
  type Rng = Xorshiftplus128Rng;

  fn _output(&self, _arm: usize) -> CommonOperatorOutput<f32> {
    assert_eq!(0, _arm);
    self.out.clone()
  }

  fn reset_loss(&mut self) {
    self.loss1 = 0.0;
    self.accuracy = 0;
  }

  fn store_loss(&mut self) -> f32 {
    self.loss1
  }

  fn forward(&mut self, _phase: OpPhase) {
    //self.out.batch_size = self.in_.batch_size;
    let batch_size = *self.in_.batch_size.borrow();
    assert_eq!(batch_size, *self.out.batch_size.borrow());

    /*for idx in 0 .. batch_size {
      let range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(self.in_.out_buf.borrow()[range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      let max_logit = self.in_.out_buf.borrow()[idx * self.cfg.num_classes + max_logit_k];
      self.max_log[idx] = max_logit;
      self.hats[idx] = max_logit_k as u32;
      for k in 0 .. self.cfg.num_classes {
        self.facts[idx * self.cfg.num_classes + k] = (self.in_.out_buf.borrow()[idx * self.cfg.num_classes + k] - max_logit).exp();
      }
      let sum_fact: f32 = Sum::sum(self.facts[range].iter().map(|&x| x));
      for k in 0 .. self.cfg.num_classes {
        self.out.out_buf.borrow_mut()[idx * self.cfg.num_classes + k] = self.facts[idx * self.cfg.num_classes + k] / sum_fact;
      }
      self.losses[idx] = -self.weights[idx] * self.out.out_buf.borrow()[idx * self.cfg.num_classes + self.labels[idx] as usize].ln();
    }*/
    let in_buf = self.in_.out_buf.borrow();
    //println!("DEBUG: softmax: in buf: {:?}", &in_buf[ .. 10]);
    let mut out_buf = self.out.out_buf.borrow_mut();
    self.sm_kern.forward(batch_size, &*in_buf, &mut *out_buf);
    let mut batch_loss = 0.0;
    let mut batch_ent = 0.0;
    let mut batch_accuracy = 0;
    for idx in 0 .. batch_size {
      let idx_range = idx * self.cfg.num_classes .. (idx+1) * self.cfg.num_classes;
      let max_logit_k = argmax(in_buf[idx_range.clone()].iter().map(|&x| F32InfNan(x))).unwrap();
      self.hats[idx] = max_logit_k as u32;
      if self.hats[idx] == self.labels[idx] {
        batch_accuracy += 1;
      }
      let (loss, ent) = if self.labels[idx] == u32::MAX {
        (0.0, 0.0)
      } else {
        let p = out_buf[idx * self.cfg.num_classes + self.labels[idx] as usize];
        let loss = -self.weights[idx] * p.ln();
        let ent = -p * p.ln();
        (loss, ent)
      };
      self.losses[idx] = loss;
      self.ents[idx] = ent;
      batch_loss += loss;
      batch_ent += ent;
    }
    //self.loss1 += Sum::sum(self.losses[ .. batch_size].iter().map(|&x| x));
    //println!("DEBUG: softmax: out buf: {:?}", &out_buf[ .. 10]);
    //println!("DEBUG: softmax: accuracy: {}/{}", batch_accuracy, batch_size);

    let in_loss = *self.in_.out_loss.borrow();
    //self.loss1 += batch_loss + in_loss;
    self.loss1 += batch_loss - self.cfg.entropy_coef * batch_ent;
    self.accuracy += batch_accuracy;
    *self.out.out_loss.borrow_mut() = batch_loss + in_loss;
  }

  fn backward(&mut self) {
    //assert_eq!(self.out.batch_size, self.in_.batch_size);
    let batch_size = *self.out.batch_size.borrow();
    if let Some(ref mut in_grad) = self.in_.out_grad.as_mut() {
      let out_buf = self.out.out_buf.borrow();
      let mut in_grad = in_grad.borrow_mut();
      let mut p = 0;
      for idx in 0 .. batch_size {
        for k in 0 .. self.cfg.num_classes {
          let y = out_buf[p];
          let w = self.weights[idx];
          let k_truth = self.labels[idx];
          let delta_k = if k == k_truth as usize { 1.0 } else { 0.0 };
          let ent = self.ents[idx];
          if y.is_nan() {
            // FIXME(20161009): try to gracefully handle NaNs?
            unimplemented!();
          } else if y == 0.0 {
            in_grad[p] = -w * delta_k;
          } else {
            in_grad[p] =
                w * (y - delta_k) + self.cfg.entropy_coef * y * (y.ln() + ent);
          }
          p += 1;
        }
      }
    }
  }
}
<|endoftext|>