<commit_before>use ide_db::imports::import_assets::item_for_path_search;
use itertools::Itertools;
use stdx::format_to;
use syntax::ast::{self, AstNode, HasName, HasVisibility, StructKind};

use crate::{
    utils::{find_impl_block_start, find_struct_impl, generate_impl_text},
    AssistContext, AssistId, AssistKind, Assists,
};

// TODO: how to depupicate with `ide-diagnostics/mssing_fields`
pub fn use_trivial_constructor(
    db: &ide_db::RootDatabase,
    path: ast::Path,
    ty: &hir::Type,
) -> Option<ast::Expr> {
    match ty.as_adt() {
        Some(hir::Adt::Enum(x)) => {
            let variants = x.variants(db);

            if variants.len() == 1 {
                let variant = variants[0];

                if variant.fields(db).is_empty() {
                    let path = ast::make::path_qualified(
                        path,
                        syntax::ast::make::path_segment(ast::make::name_ref(
                            &variant.name(db).to_smol_str(),
                        )),
                    );

                    let is_record = variant.kind(db) == hir::StructKind::Record;

                    return Some(if is_record {
                        ast::Expr::RecordExpr(syntax::ast::make::record_expr(
                            path,
                            ast::make::record_expr_field_list(std::iter::empty()),
                        ))
                    } else {
                        syntax::ast::make::expr_path(path)
                    });
                }
            }
        }
        Some(hir::Adt::Struct(x)) => {
            if x.fields(db).is_empty() {
                return Some(syntax::ast::make::expr_path(path));
            }
        }
        _ => {}
    }

    None
}

// Assist: generate_new
//
// Adds a new inherent impl for a type.
//
// ```
// struct Ctx<T: Clone> {
//      data: T,$0
// }
// ```
// ->
// ```
// struct Ctx<T: Clone> {
//      data: T,
// }
//
// impl<T: Clone> Ctx<T> {
//     fn $0new(data: T) -> Self { Self { data } }
// }
// ```
pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let strukt = ctx.find_node_at_offset::<ast::Struct>()?;

    // We want to only apply this to non-union structs with named fields
    let field_list = match strukt.kind() {
        StructKind::Record(named) => named,
        _ => return None,
    };

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), "new")?;

    let target = strukt.syntax().text_range();
    acc.add(AssistId("generate_new", AssistKind::Generate), "Generate `new`", target, |builder| {
        let mut buf = String::with_capacity(512);

        if impl_def.is_some() {
            buf.push('\n');
        }

        let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));

        let current_module = ctx.sema.scope(strukt.syntax()).unwrap().module();

        let trivial_constructors = field_list
            .fields()
            .map(|f| {
                let ty = ctx.sema.resolve_type(&f.ty()?)?;

                let item_in_ns = hir::ItemInNs::from(hir::ModuleDef::from(ty.as_adt()?));

                let type_path = current_module
                    .find_use_path(ctx.sema.db, item_for_path_search(ctx.sema.db, item_in_ns)?)?;

                let expr = use_trivial_constructor(
                    &ctx.sema.db,
                    ide_db::helpers::mod_path_to_ast(&type_path),
                    &ty,
                )?;

                Some(format!("{}: {}", f.name()?.syntax(), expr))
            })
            .collect::<Vec<_>>();

        dbg!(&trivial_constructors);

        let params = field_list
            .fields()
            .enumerate()
            .filter_map(|(i, f)| {
                if trivial_constructors[i].is_none() {
                    Some(format!("{}: {}", f.name()?.syntax(), f.ty()?.syntax()))
                } else {
                    None
                }
            })
            .format(", ");

        let fields = field_list
            .fields()
            .enumerate()
            .filter_map(|(i, f)| {
                let contructor = trivial_constructors[i].clone();
                if contructor.is_some() {
                    contructor
                } else {
                    Some(f.name()?.to_string())
                }
            })
            .format(", ");

        format_to!(buf, "    {}fn new({}) -> Self {{ Self {{ {} }} }}", vis, params, fields);

        let start_offset = impl_def
            .and_then(|impl_def| find_impl_block_start(impl_def, &mut buf))
            .unwrap_or_else(|| {
                buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
                strukt.syntax().text_range().end()
            });

        match ctx.config.snippet_cap {
            None => builder.insert(start_offset, buf),
            Some(cap) => {
                buf = buf.replace("fn new", "fn $0new");
                builder.insert_snippet(cap, start_offset, buf);
            }
        }
    })
}

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

    use super::*;

    #[test]
    fn test_generate_new() {
        check_assist(
            generate_new,
            r#"
struct Foo {$0}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo<T: Clone> {$0}
"#,
            r#"
struct Foo<T: Clone> {}

impl<T: Clone> Foo<T> {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo<'a, T: Foo<'a>> {$0}
"#,
            r#"
struct Foo<'a, T: Foo<'a>> {}

impl<'a, T: Foo<'a>> Foo<'a, T> {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo { baz: String $0}
"#,
            r#"
struct Foo { baz: String }

impl Foo {
    fn $0new(baz: String) -> Self { Self { baz } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo { baz: String, qux: Vec<i32> $0}
"#,
            r#"
struct Foo { baz: String, qux: Vec<i32> }

impl Foo {
    fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
}
"#,
        );
    }

    #[test]
    fn check_that_visibility_modifiers_dont_get_brought_in() {
        check_assist(
            generate_new,
            r#"
struct Foo { pub baz: String, pub qux: Vec<i32> $0}
"#,
            r#"
struct Foo { pub baz: String, pub qux: Vec<i32> }

impl Foo {
    fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
}
"#,
        );
    }

    #[test]
    fn check_it_reuses_existing_impls() {
        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn qux(&self) {}
}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }

    fn qux(&self) {}
}
"#,
        );

        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }

    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
"#,
        );
    }

    #[test]
    fn check_visibility_of_new_fn_based_on_struct() {
        check_assist(
            generate_new,
            r#"
pub struct Foo {$0}
"#,
            r#"
pub struct Foo {}

impl Foo {
    pub fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
pub(crate) struct Foo {$0}
"#,
            r#"
pub(crate) struct Foo {}

impl Foo {
    pub(crate) fn $0new() -> Self { Self {  } }
}
"#,
        );
    }

    #[test]
    fn generate_new_not_applicable_if_fn_exists() {
        check_assist_not_applicable(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn new() -> Self {
        Self
    }
}
"#,
        );

        check_assist_not_applicable(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn New() -> Self {
        Self
    }
}
"#,
        );
    }

    #[test]
    fn generate_new_target() {
        check_assist_target(
            generate_new,
            r#"
struct SomeThingIrrelevant;
/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {$0}
struct EvenMoreIrrelevant;
"#,
            "/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {}",
        );
    }

    #[test]
    fn test_unrelated_new() {
        check_assist(
            generate_new,
            r#"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,$0
    pub ast: T,
}

impl<T> Source<T> {
    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"#,
            r#"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,
    pub ast: T,
}

impl<T> Source<T> {
    pub fn $0new(file_id: HirFileId, ast: T) -> Self { Self { file_id, ast } }

    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"#,
        );
    }
}
<commit_msg>remove dbg<commit_after>use ide_db::imports::import_assets::item_for_path_search;
use itertools::Itertools;
use stdx::format_to;
use syntax::ast::{self, AstNode, HasName, HasVisibility, StructKind};

use crate::{
    utils::{find_impl_block_start, find_struct_impl, generate_impl_text},
    AssistContext, AssistId, AssistKind, Assists,
};

// TODO: how to depupicate with `ide-diagnostics/mssing_fields`
pub fn use_trivial_constructor(
    db: &ide_db::RootDatabase,
    path: ast::Path,
    ty: &hir::Type,
) -> Option<ast::Expr> {
    match ty.as_adt() {
        Some(hir::Adt::Enum(x)) => {
            let variants = x.variants(db);

            if variants.len() == 1 {
                let variant = variants[0];

                if variant.fields(db).is_empty() {
                    let path = ast::make::path_qualified(
                        path,
                        syntax::ast::make::path_segment(ast::make::name_ref(
                            &variant.name(db).to_smol_str(),
                        )),
                    );

                    let is_record = variant.kind(db) == hir::StructKind::Record;

                    return Some(if is_record {
                        ast::Expr::RecordExpr(syntax::ast::make::record_expr(
                            path,
                            ast::make::record_expr_field_list(std::iter::empty()),
                        ))
                    } else {
                        syntax::ast::make::expr_path(path)
                    });
                }
            }
        }
        Some(hir::Adt::Struct(x)) => {
            if x.fields(db).is_empty() {
                return Some(syntax::ast::make::expr_path(path));
            }
        }
        _ => {}
    }

    None
}

// Assist: generate_new
//
// Adds a new inherent impl for a type.
//
// ```
// struct Ctx<T: Clone> {
//      data: T,$0
// }
// ```
// ->
// ```
// struct Ctx<T: Clone> {
//      data: T,
// }
//
// impl<T: Clone> Ctx<T> {
//     fn $0new(data: T) -> Self { Self { data } }
// }
// ```
pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let strukt = ctx.find_node_at_offset::<ast::Struct>()?;

    // We want to only apply this to non-union structs with named fields
    let field_list = match strukt.kind() {
        StructKind::Record(named) => named,
        _ => return None,
    };

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), "new")?;

    let target = strukt.syntax().text_range();
    acc.add(AssistId("generate_new", AssistKind::Generate), "Generate `new`", target, |builder| {
        let mut buf = String::with_capacity(512);

        if impl_def.is_some() {
            buf.push('\n');
        }

        let vis = strukt.visibility().map_or(String::new(), |v| format!("{} ", v));

        let current_module = ctx.sema.scope(strukt.syntax()).unwrap().module();

        let trivial_constructors = field_list
            .fields()
            .map(|f| {
                let ty = ctx.sema.resolve_type(&f.ty()?)?;

                let item_in_ns = hir::ItemInNs::from(hir::ModuleDef::from(ty.as_adt()?));

                let type_path = current_module
                    .find_use_path(ctx.sema.db, item_for_path_search(ctx.sema.db, item_in_ns)?)?;

                let expr = use_trivial_constructor(
                    &ctx.sema.db,
                    ide_db::helpers::mod_path_to_ast(&type_path),
                    &ty,
                )?;

                Some(format!("{}: {}", f.name()?.syntax(), expr))
            })
            .collect::<Vec<_>>();

        let params = field_list
            .fields()
            .enumerate()
            .filter_map(|(i, f)| {
                if trivial_constructors[i].is_none() {
                    Some(format!("{}: {}", f.name()?.syntax(), f.ty()?.syntax()))
                } else {
                    None
                }
            })
            .format(", ");

        let fields = field_list
            .fields()
            .enumerate()
            .filter_map(|(i, f)| {
                let contructor = trivial_constructors[i].clone();
                if contructor.is_some() {
                    contructor
                } else {
                    Some(f.name()?.to_string())
                }
            })
            .format(", ");

        format_to!(buf, "    {}fn new({}) -> Self {{ Self {{ {} }} }}", vis, params, fields);

        let start_offset = impl_def
            .and_then(|impl_def| find_impl_block_start(impl_def, &mut buf))
            .unwrap_or_else(|| {
                buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
                strukt.syntax().text_range().end()
            });

        match ctx.config.snippet_cap {
            None => builder.insert(start_offset, buf),
            Some(cap) => {
                buf = buf.replace("fn new", "fn $0new");
                builder.insert_snippet(cap, start_offset, buf);
            }
        }
    })
}

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

    use super::*;

    #[test]
    fn test_generate_new() {
        check_assist(
            generate_new,
            r#"
struct Foo {$0}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo<T: Clone> {$0}
"#,
            r#"
struct Foo<T: Clone> {}

impl<T: Clone> Foo<T> {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo<'a, T: Foo<'a>> {$0}
"#,
            r#"
struct Foo<'a, T: Foo<'a>> {}

impl<'a, T: Foo<'a>> Foo<'a, T> {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo { baz: String $0}
"#,
            r#"
struct Foo { baz: String }

impl Foo {
    fn $0new(baz: String) -> Self { Self { baz } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo { baz: String, qux: Vec<i32> $0}
"#,
            r#"
struct Foo { baz: String, qux: Vec<i32> }

impl Foo {
    fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
}
"#,
        );
    }

    #[test]
    fn check_that_visibility_modifiers_dont_get_brought_in() {
        check_assist(
            generate_new,
            r#"
struct Foo { pub baz: String, pub qux: Vec<i32> $0}
"#,
            r#"
struct Foo { pub baz: String, pub qux: Vec<i32> }

impl Foo {
    fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
}
"#,
        );
    }

    #[test]
    fn check_it_reuses_existing_impls() {
        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn qux(&self) {}
}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }

    fn qux(&self) {}
}
"#,
        );

        check_assist(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
"#,
            r#"
struct Foo {}

impl Foo {
    fn $0new() -> Self { Self {  } }

    fn qux(&self) {}
    fn baz() -> i32 {
        5
    }
}
"#,
        );
    }

    #[test]
    fn check_visibility_of_new_fn_based_on_struct() {
        check_assist(
            generate_new,
            r#"
pub struct Foo {$0}
"#,
            r#"
pub struct Foo {}

impl Foo {
    pub fn $0new() -> Self { Self {  } }
}
"#,
        );
        check_assist(
            generate_new,
            r#"
pub(crate) struct Foo {$0}
"#,
            r#"
pub(crate) struct Foo {}

impl Foo {
    pub(crate) fn $0new() -> Self { Self {  } }
}
"#,
        );
    }

    #[test]
    fn generate_new_not_applicable_if_fn_exists() {
        check_assist_not_applicable(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn new() -> Self {
        Self
    }
}
"#,
        );

        check_assist_not_applicable(
            generate_new,
            r#"
struct Foo {$0}

impl Foo {
    fn New() -> Self {
        Self
    }
}
"#,
        );
    }

    #[test]
    fn generate_new_target() {
        check_assist_target(
            generate_new,
            r#"
struct SomeThingIrrelevant;
/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {$0}
struct EvenMoreIrrelevant;
"#,
            "/// Has a lifetime parameter
struct Foo<'a, T: Foo<'a>> {}",
        );
    }

    #[test]
    fn test_unrelated_new() {
        check_assist(
            generate_new,
            r#"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,$0
    pub ast: T,
}

impl<T> Source<T> {
    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"#,
            r#"
pub struct AstId<N: AstNode> {
    file_id: HirFileId,
    file_ast_id: FileAstId<N>,
}

impl<N: AstNode> AstId<N> {
    pub fn new(file_id: HirFileId, file_ast_id: FileAstId<N>) -> AstId<N> {
        AstId { file_id, file_ast_id }
    }
}

pub struct Source<T> {
    pub file_id: HirFileId,
    pub ast: T,
}

impl<T> Source<T> {
    pub fn $0new(file_id: HirFileId, ast: T) -> Self { Self { file_id, ast } }

    pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
        Source { file_id: self.file_id, ast: f(self.ast) }
    }
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Mozilla
//
// 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.

extern crate ordered_float;

use std::collections::BTreeMap;
use self::ordered_float::OrderedFloat;

/// Core types defining a Mentat knowledge base.

/// Represents one entid in the entid space.
///
/// Per https://www.sqlite.org/datatype3.html (see also http://stackoverflow.com/a/8499544), SQLite
/// stores signed integers up to 64 bits in size.  Since u32 is not appropriate for our use case, we
/// use i64 rather than manually truncating u64 to u63 and casting to i64 throughout the codebase.
pub type Entid = i64;

/// The attribute of each Mentat assertion has a :db/valueType constraining the value to a
/// particular set.  Mentat recognizes the following :db/valueType values.
#[derive(Clone,Copy,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub enum ValueType {
    Ref,
    Boolean,
    Instant,
    Long,
    Double,
    String,
    Keyword,
}

/// Represents a Mentat value in a particular value set.
// TODO: expand to include :db.type/{instant,url,uuid}.
// TODO: BigInt?
#[derive(Clone,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub enum TypedValue {
    Ref(Entid),
    Boolean(bool),
    Long(i64),
    Double(OrderedFloat<f64>),
    // TODO: &str throughout?
    String(String),
    Keyword(String),
}

impl TypedValue {
    /// Returns true if the provided type is `Some` and matches this value's type, or if the
    /// provided type is `None`.
    #[inline]
    pub fn is_congruent_with<T: Into<Option<ValueType>>>(&self, t: T) -> bool {
        t.into().map_or(true, |x| self.matches_type(x))
    }

    #[inline]
    pub fn matches_type(&self, t: ValueType) -> bool {
        self.value_type() == t
    }

    pub fn value_type(&self) -> ValueType {
        match self {
            &TypedValue::Ref(_) => ValueType::Ref,
            &TypedValue::Boolean(_) => ValueType::Boolean,
            &TypedValue::Long(_) => ValueType::Long,
            &TypedValue::Double(_) => ValueType::Double,
            &TypedValue::String(_) => ValueType::String,
            &TypedValue::Keyword(_) => ValueType::Keyword,
        }
    }
}

#[test]
fn test_typed_value() {
    assert!(TypedValue::Boolean(false).is_congruent_with(None));
    assert!(TypedValue::Boolean(false).is_congruent_with(ValueType::Boolean));
    assert!(!TypedValue::String("foo".to_string()).is_congruent_with(ValueType::Boolean));
    assert!(TypedValue::String("foo".to_string()).is_congruent_with(ValueType::String));
    assert!(TypedValue::String("foo".to_string()).is_congruent_with(None));
}

/// Bit flags used in `flags0` column in temporary tables created during search,
/// such as the `search_results`, `inexact_searches` and `exact_searches` tables.
/// When moving to a more concrete table, such as `datoms`, they are expanded out
/// via these flags and put into their own column rather than a bit field.
pub enum AttributeBitFlags {
    IndexAVET     = 1 << 0,
    IndexVAET     = 1 << 1,
    IndexFulltext = 1 << 2,
    UniqueValue   = 1 << 3,
}

/// A Mentat schema attribute has a value type and several other flags determining how assertions
/// with the attribute are interpreted.
///
/// TODO: consider packing this into a bitfield or similar.
#[derive(Clone,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub struct Attribute {
    /// The associated value type, i.e., `:db/valueType`?
    pub value_type: ValueType,

    /// `true` if this attribute is multi-valued, i.e., it is `:db/cardinality
    /// :db.cardinality/many`.  `false` if this attribute is single-valued (the default), i.e., it
    /// is `:db/cardinality :db.cardinality/one`.
    pub multival: bool,

    /// `true` if this attribute is unique-value, i.e., it is `:db/unique :db.unique/value`.
    ///
    /// *Unique-value* means that there is at most one assertion with the attribute and a
    /// particular value in the datom store.
    pub unique_value: bool,

    /// `true` if this attribute is unique-identity, i.e., it is `:db/unique :db.unique/identity`.
    ///
    /// Unique-identity attributes always have value type `Ref`.
    ///
    /// *Unique-identity* means that the attribute is *unique-value* and that they can be used in
    /// lookup-refs and will automatically upsert where appropriate.
    pub unique_identity: bool,

    /// `true` if this attribute is automatically indexed, i.e., it is `:db/indexing true`.
    pub index: bool,

    /// `true` if this attribute is automatically fulltext indexed, i.e., it is `:db/fulltext true`.
    ///
    /// Fulltext attributes always have string values.
    pub fulltext: bool,

    /// `true` if this attribute is a component, i.e., it is `:db/isComponent true`.
    ///
    /// Component attributes always have value type `Ref`.
    ///
    /// They are used to compose entities from component sub-entities: they are fetched recursively
    /// by pull expressions, and they are automatically recursively deleted where appropriate.
    pub component: bool,
}

impl Attribute {
    /// Combine several attribute flags into a bitfield used in temporary search tables.
    pub fn flags(&self) -> u8 {
        let mut flags: u8 = 0;

        if self.index {
            flags |= AttributeBitFlags::IndexAVET as u8;
        }
        if self.value_type == ValueType::Ref {
            flags |= AttributeBitFlags::IndexVAET as u8;
        }
        if self.fulltext {
            flags |= AttributeBitFlags::IndexFulltext as u8;
        }
        if self.unique_value {
            flags |= AttributeBitFlags::UniqueValue as u8;
        }
        flags
    }
}

impl Default for Attribute {
    fn default() -> Attribute {
        Attribute {
            // There's no particular reason to favour one value type, so Ref it is.
            value_type: ValueType::Ref,
            fulltext: false,
            index: false,
            multival: false,
            unique_value: false,
            unique_identity: false,
            component: false,
        }
    }
}

/// Map `String` idents (`:db/ident`) to positive integer entids (`1`).
/// TODO: these should all be parsed into NamespacedKeywords on entry. #291.
pub type IdentMap = BTreeMap<String, Entid>;

/// Map positive integer entids (`1`) to `String` idents (`:db/ident`).
pub type EntidMap = BTreeMap<Entid, String>;

/// Map attribute entids to `Attribute` instances.
pub type SchemaMap = BTreeMap<Entid, Attribute>;

/// Represents a Mentat schema.
///
/// Maintains the mapping between string idents and positive integer entids; and exposes the schema
/// flags associated to a given entid (equivalently, ident).
///
/// TODO: consider a single bi-directional map instead of separate ident->entid and entid->ident
/// maps.
#[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialOrd, PartialEq)]
pub struct Schema {
    /// Map entid->ident.
    ///
    /// Invariant: is the inverse map of `ident_map`.
    pub entid_map: EntidMap,

    /// Map ident->entid.
    ///
    /// Invariant: is the inverse map of `entid_map`.
    pub ident_map: IdentMap,

    /// Map entid->attribute flags.
    ///
    /// Invariant: key-set is the same as the key-set of `entid_map` (equivalently, the value-set of
    /// `ident_map`).
    pub schema_map: SchemaMap,
}

impl Schema {
    pub fn get_ident(&self, x: Entid) -> Option<&String> {
        self.entid_map.get(&x)
    }

    pub fn get_entid(&self, x: &String) -> Option<Entid> {
        self.ident_map.get(x).map(|x| *x)
    }

    pub fn attribute_for_entid(&self, x: Entid) -> Option<&Attribute> {
        self.schema_map.get(&x)
    }

    pub fn is_attribute(&self, x: Entid) -> bool {
        self.schema_map.contains_key(&x)
    }

    pub fn identifies_attribute(&self, x: &String) -> bool {
        self.get_entid(x).map(|e| self.is_attribute(e)).unwrap_or(false)
    }
}

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

    #[test]
    fn test_attribute_flags() {
        let attr1 = Attribute {
            index: true,
            value_type: ValueType::Ref,
            fulltext: false,
            unique_value: false,
            multival: false,
            unique_identity: false,
            component: false,
        };

        assert!(attr1.flags() & AttributeBitFlags::IndexAVET as u8 != 0);
        assert!(attr1.flags() & AttributeBitFlags::IndexVAET as u8 != 0);
        assert!(attr1.flags() & AttributeBitFlags::IndexFulltext as u8 == 0);
        assert!(attr1.flags() & AttributeBitFlags::UniqueValue as u8 == 0);

        let attr2 = Attribute {
            index: false,
            value_type: ValueType::Boolean,
            fulltext: true,
            unique_value: true,
            multival: false,
            unique_identity: false,
            component: false,
        };

        assert!(attr2.flags() & AttributeBitFlags::IndexAVET as u8 == 0);
        assert!(attr2.flags() & AttributeBitFlags::IndexVAET as u8 == 0);
        assert!(attr2.flags() & AttributeBitFlags::IndexFulltext as u8 != 0);
        assert!(attr2.flags() & AttributeBitFlags::UniqueValue as u8 != 0);
    }
}
<commit_msg>Add docstrings for Schema methods.<commit_after>// Copyright 2016 Mozilla
//
// 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.

extern crate ordered_float;

use std::collections::BTreeMap;
use self::ordered_float::OrderedFloat;

/// Core types defining a Mentat knowledge base.

/// Represents one entid in the entid space.
///
/// Per https://www.sqlite.org/datatype3.html (see also http://stackoverflow.com/a/8499544), SQLite
/// stores signed integers up to 64 bits in size.  Since u32 is not appropriate for our use case, we
/// use i64 rather than manually truncating u64 to u63 and casting to i64 throughout the codebase.
pub type Entid = i64;

/// The attribute of each Mentat assertion has a :db/valueType constraining the value to a
/// particular set.  Mentat recognizes the following :db/valueType values.
#[derive(Clone,Copy,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub enum ValueType {
    Ref,
    Boolean,
    Instant,
    Long,
    Double,
    String,
    Keyword,
}

/// Represents a Mentat value in a particular value set.
// TODO: expand to include :db.type/{instant,url,uuid}.
// TODO: BigInt?
#[derive(Clone,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub enum TypedValue {
    Ref(Entid),
    Boolean(bool),
    Long(i64),
    Double(OrderedFloat<f64>),
    // TODO: &str throughout?
    String(String),
    Keyword(String),
}

impl TypedValue {
    /// Returns true if the provided type is `Some` and matches this value's type, or if the
    /// provided type is `None`.
    #[inline]
    pub fn is_congruent_with<T: Into<Option<ValueType>>>(&self, t: T) -> bool {
        t.into().map_or(true, |x| self.matches_type(x))
    }

    #[inline]
    pub fn matches_type(&self, t: ValueType) -> bool {
        self.value_type() == t
    }

    pub fn value_type(&self) -> ValueType {
        match self {
            &TypedValue::Ref(_) => ValueType::Ref,
            &TypedValue::Boolean(_) => ValueType::Boolean,
            &TypedValue::Long(_) => ValueType::Long,
            &TypedValue::Double(_) => ValueType::Double,
            &TypedValue::String(_) => ValueType::String,
            &TypedValue::Keyword(_) => ValueType::Keyword,
        }
    }
}

#[test]
fn test_typed_value() {
    assert!(TypedValue::Boolean(false).is_congruent_with(None));
    assert!(TypedValue::Boolean(false).is_congruent_with(ValueType::Boolean));
    assert!(!TypedValue::String("foo".to_string()).is_congruent_with(ValueType::Boolean));
    assert!(TypedValue::String("foo".to_string()).is_congruent_with(ValueType::String));
    assert!(TypedValue::String("foo".to_string()).is_congruent_with(None));
}

/// Bit flags used in `flags0` column in temporary tables created during search,
/// such as the `search_results`, `inexact_searches` and `exact_searches` tables.
/// When moving to a more concrete table, such as `datoms`, they are expanded out
/// via these flags and put into their own column rather than a bit field.
pub enum AttributeBitFlags {
    IndexAVET     = 1 << 0,
    IndexVAET     = 1 << 1,
    IndexFulltext = 1 << 2,
    UniqueValue   = 1 << 3,
}

/// A Mentat schema attribute has a value type and several other flags determining how assertions
/// with the attribute are interpreted.
///
/// TODO: consider packing this into a bitfield or similar.
#[derive(Clone,Debug,Eq,Hash,Ord,PartialOrd,PartialEq)]
pub struct Attribute {
    /// The associated value type, i.e., `:db/valueType`?
    pub value_type: ValueType,

    /// `true` if this attribute is multi-valued, i.e., it is `:db/cardinality
    /// :db.cardinality/many`.  `false` if this attribute is single-valued (the default), i.e., it
    /// is `:db/cardinality :db.cardinality/one`.
    pub multival: bool,

    /// `true` if this attribute is unique-value, i.e., it is `:db/unique :db.unique/value`.
    ///
    /// *Unique-value* means that there is at most one assertion with the attribute and a
    /// particular value in the datom store.
    pub unique_value: bool,

    /// `true` if this attribute is unique-identity, i.e., it is `:db/unique :db.unique/identity`.
    ///
    /// Unique-identity attributes always have value type `Ref`.
    ///
    /// *Unique-identity* means that the attribute is *unique-value* and that they can be used in
    /// lookup-refs and will automatically upsert where appropriate.
    pub unique_identity: bool,

    /// `true` if this attribute is automatically indexed, i.e., it is `:db/indexing true`.
    pub index: bool,

    /// `true` if this attribute is automatically fulltext indexed, i.e., it is `:db/fulltext true`.
    ///
    /// Fulltext attributes always have string values.
    pub fulltext: bool,

    /// `true` if this attribute is a component, i.e., it is `:db/isComponent true`.
    ///
    /// Component attributes always have value type `Ref`.
    ///
    /// They are used to compose entities from component sub-entities: they are fetched recursively
    /// by pull expressions, and they are automatically recursively deleted where appropriate.
    pub component: bool,
}

impl Attribute {
    /// Combine several attribute flags into a bitfield used in temporary search tables.
    pub fn flags(&self) -> u8 {
        let mut flags: u8 = 0;

        if self.index {
            flags |= AttributeBitFlags::IndexAVET as u8;
        }
        if self.value_type == ValueType::Ref {
            flags |= AttributeBitFlags::IndexVAET as u8;
        }
        if self.fulltext {
            flags |= AttributeBitFlags::IndexFulltext as u8;
        }
        if self.unique_value {
            flags |= AttributeBitFlags::UniqueValue as u8;
        }
        flags
    }
}

impl Default for Attribute {
    fn default() -> Attribute {
        Attribute {
            // There's no particular reason to favour one value type, so Ref it is.
            value_type: ValueType::Ref,
            fulltext: false,
            index: false,
            multival: false,
            unique_value: false,
            unique_identity: false,
            component: false,
        }
    }
}

/// Map `String` idents (`:db/ident`) to positive integer entids (`1`).
/// TODO: these should all be parsed into NamespacedKeywords on entry. #291.
pub type IdentMap = BTreeMap<String, Entid>;

/// Map positive integer entids (`1`) to `String` idents (`:db/ident`).
pub type EntidMap = BTreeMap<Entid, String>;

/// Map attribute entids to `Attribute` instances.
pub type SchemaMap = BTreeMap<Entid, Attribute>;

/// Represents a Mentat schema.
///
/// Maintains the mapping between string idents and positive integer entids; and exposes the schema
/// flags associated to a given entid (equivalently, ident).
///
/// TODO: consider a single bi-directional map instead of separate ident->entid and entid->ident
/// maps.
#[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialOrd, PartialEq)]
pub struct Schema {
    /// Map entid->ident.
    ///
    /// Invariant: is the inverse map of `ident_map`.
    pub entid_map: EntidMap,

    /// Map ident->entid.
    ///
    /// Invariant: is the inverse map of `entid_map`.
    pub ident_map: IdentMap,

    /// Map entid->attribute flags.
    ///
    /// Invariant: key-set is the same as the key-set of `entid_map` (equivalently, the value-set of
    /// `ident_map`).
    pub schema_map: SchemaMap,
}

impl Schema {
    pub fn get_ident(&self, x: Entid) -> Option<&String> {
        self.entid_map.get(&x)
    }

    pub fn get_entid(&self, x: &String) -> Option<Entid> {
        self.ident_map.get(x).map(|x| *x)
    }

    pub fn attribute_for_entid(&self, x: Entid) -> Option<&Attribute> {
        self.schema_map.get(&x)
    }

    /// Return true if the provided entid identifies an attribute in this schema.
    pub fn is_attribute(&self, x: Entid) -> bool {
        self.schema_map.contains_key(&x)
    }

    /// Return true if the provided ident identifies an attribute in this schema.
    pub fn identifies_attribute(&self, x: &String) -> bool {
        self.get_entid(x).map(|e| self.is_attribute(e)).unwrap_or(false)
    }
}

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

    #[test]
    fn test_attribute_flags() {
        let attr1 = Attribute {
            index: true,
            value_type: ValueType::Ref,
            fulltext: false,
            unique_value: false,
            multival: false,
            unique_identity: false,
            component: false,
        };

        assert!(attr1.flags() & AttributeBitFlags::IndexAVET as u8 != 0);
        assert!(attr1.flags() & AttributeBitFlags::IndexVAET as u8 != 0);
        assert!(attr1.flags() & AttributeBitFlags::IndexFulltext as u8 == 0);
        assert!(attr1.flags() & AttributeBitFlags::UniqueValue as u8 == 0);

        let attr2 = Attribute {
            index: false,
            value_type: ValueType::Boolean,
            fulltext: true,
            unique_value: true,
            multival: false,
            unique_identity: false,
            component: false,
        };

        assert!(attr2.flags() & AttributeBitFlags::IndexAVET as u8 == 0);
        assert!(attr2.flags() & AttributeBitFlags::IndexVAET as u8 == 0);
        assert!(attr2.flags() & AttributeBitFlags::IndexFulltext as u8 != 0);
        assert!(attr2.flags() & AttributeBitFlags::UniqueValue as u8 != 0);
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// 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.

//! The proof of work needs to strike a balance between fast header
//! verification to avoid DoS attacks and difficulty for block verifiers to
//! build new blocks. In addition, mining new blocks should also be as
//! difficult on high end custom-made hardware (ASICs) as on commodity hardware
//! or smartphones. For this reason we use Cuckoo Cycle (see the cuckoo
//! module for more information).
//!
//! Note that this miner implementation is here mostly for tests and
//! reference. It's not optimized for speed.

#![deny(non_upper_case_globals)]
#![deny(non_camel_case_types)]
#![deny(non_snake_case)]
#![deny(unused_mut)]
#![warn(missing_docs)]

use chrono;
use num;

#[macro_use]
mod common;
pub mod cuckaroo;
pub mod cuckatoo;
mod error;
#[allow(dead_code)]
pub mod lean;
mod siphash;
mod types;

use crate::core::{Block, BlockHeader};
use crate::genesis;
use crate::global;
use chrono::prelude::{DateTime, NaiveDateTime, Utc};

pub use self::common::EdgeType;
pub use self::types::*;
pub use crate::pow::cuckaroo::{new_cuckaroo_ctx, CuckarooContext};
pub use crate::pow::cuckatoo::{new_cuckatoo_ctx, CuckatooContext};
pub use crate::pow::error::Error;

const MAX_SOLS: u32 = 10;

/// Validates the proof of work of a given header, and that the proof of work
/// satisfies the requirements of the header.
pub fn verify_size(bh: &BlockHeader) -> Result<(), Error> {
	let mut ctx = global::create_pow_context::<u64>(
		bh.height,
		bh.pow.edge_bits(),
		bh.pow.proof.nonces.len(),
		MAX_SOLS,
	)?;
	ctx.set_header_nonce(bh.pre_pow(), None, false)?;
	ctx.verify(&bh.pow.proof)
}

/// Mines a genesis block using the internal miner
pub fn mine_genesis_block() -> Result<Block, Error> {
	let mut gen = genesis::genesis_dev();
	if global::is_user_testing_mode() || global::is_automated_testing_mode() {
		gen = genesis::genesis_dev();
		gen.header.timestamp = Utc::now();
	}

	// total_difficulty on the genesis header *is* the difficulty of that block
	let genesis_difficulty = gen.header.pow.total_difficulty;

	let sz = global::min_edge_bits();
	let proof_size = global::proofsize();

	pow_size(&mut gen.header, genesis_difficulty, proof_size, sz)?;
	Ok(gen)
}

/// Runs a proof of work computation over the provided block using the provided
/// Mining Worker, until the required difficulty target is reached. May take a
/// while for a low target...
pub fn pow_size(
	bh: &mut BlockHeader,
	diff: Difficulty,
	proof_size: usize,
	sz: u8,
) -> Result<(), Error> {
	let start_nonce = bh.pow.nonce;

	// set the nonce for faster solution finding in user testing
	if bh.height == 0 && global::is_user_testing_mode() {
		bh.pow.nonce = global::get_genesis_nonce();
	}

	// try to find a cuckoo cycle on that header hash
	loop {
		// if we found a cycle (not guaranteed) and the proof hash is higher that the
		// diff, we're all good
		let mut ctx = global::create_pow_context::<u32>(bh.height, sz, proof_size, MAX_SOLS)?;
		ctx.set_header_nonce(bh.pre_pow(), None, true)?;
		if let Ok(proofs) = ctx.find_cycles() {
			bh.pow.proof = proofs[0].clone();
			if bh.pow.to_difficulty(bh.height) >= diff {
				return Ok(());
			}
		}

		// otherwise increment the nonce
		let (res, _) = bh.pow.nonce.overflowing_add(1);
		bh.pow.nonce = res;

		// and if we're back where we started, update the time (changes the hash as
		// well)
		if bh.pow.nonce == start_nonce {
			bh.timestamp = DateTime::<Utc>::from_utc(NaiveDateTime::from_timestamp(0, 0), Utc);
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use crate::genesis;
	use crate::global;
	use crate::global::ChainTypes;

	/// We'll be generating genesis blocks differently
	#[test]
	fn genesis_pow() {
		global::set_mining_mode(ChainTypes::UserTesting);

		let mut b = genesis::genesis_dev();
		b.header.pow.nonce = 28106;
		b.header.pow.proof.edge_bits = global::min_edge_bits();
		println!("proof {}", global::proofsize());
		pow_size(
			&mut b.header,
			Difficulty::min(),
			global::proofsize(),
			global::min_edge_bits(),
		)
		.unwrap();
		println!("nonce {}", b.header.pow.nonce);
		assert_ne!(b.header.pow.nonce, 310);
		assert!(b.header.pow.to_difficulty(0) >= Difficulty::min());
		let start = ::std::time::Instant::now();
		for n in 0..100000 {
			assert!(verify_size(&b.header).is_ok());
		}
		println!("==> {}", start.elapsed().as_secs());
	}
}
<commit_msg>Test cleanup<commit_after>// Copyright 2018 The Grin Developers
//
// 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.

//! The proof of work needs to strike a balance between fast header
//! verification to avoid DoS attacks and difficulty for block verifiers to
//! build new blocks. In addition, mining new blocks should also be as
//! difficult on high end custom-made hardware (ASICs) as on commodity hardware
//! or smartphones. For this reason we use Cuckoo Cycle (see the cuckoo
//! module for more information).
//!
//! Note that this miner implementation is here mostly for tests and
//! reference. It's not optimized for speed.

#![deny(non_upper_case_globals)]
#![deny(non_camel_case_types)]
#![deny(non_snake_case)]
#![deny(unused_mut)]
#![warn(missing_docs)]

use chrono;
use num;

#[macro_use]
mod common;
pub mod cuckaroo;
pub mod cuckatoo;
mod error;
#[allow(dead_code)]
pub mod lean;
mod siphash;
mod types;

use crate::core::{Block, BlockHeader};
use crate::genesis;
use crate::global;
use chrono::prelude::{DateTime, NaiveDateTime, Utc};

pub use self::common::EdgeType;
pub use self::types::*;
pub use crate::pow::cuckaroo::{new_cuckaroo_ctx, CuckarooContext};
pub use crate::pow::cuckatoo::{new_cuckatoo_ctx, CuckatooContext};
pub use crate::pow::error::Error;

const MAX_SOLS: u32 = 10;

/// Validates the proof of work of a given header, and that the proof of work
/// satisfies the requirements of the header.
pub fn verify_size(bh: &BlockHeader) -> Result<(), Error> {
	let mut ctx = global::create_pow_context::<u64>(
		bh.height,
		bh.pow.edge_bits(),
		bh.pow.proof.nonces.len(),
		MAX_SOLS,
	)?;
	ctx.set_header_nonce(bh.pre_pow(), None, false)?;
	ctx.verify(&bh.pow.proof)
}

/// Mines a genesis block using the internal miner
pub fn mine_genesis_block() -> Result<Block, Error> {
	let mut gen = genesis::genesis_dev();
	if global::is_user_testing_mode() || global::is_automated_testing_mode() {
		gen = genesis::genesis_dev();
		gen.header.timestamp = Utc::now();
	}

	// total_difficulty on the genesis header *is* the difficulty of that block
	let genesis_difficulty = gen.header.pow.total_difficulty;

	let sz = global::min_edge_bits();
	let proof_size = global::proofsize();

	pow_size(&mut gen.header, genesis_difficulty, proof_size, sz)?;
	Ok(gen)
}

/// Runs a proof of work computation over the provided block using the provided
/// Mining Worker, until the required difficulty target is reached. May take a
/// while for a low target...
pub fn pow_size(
	bh: &mut BlockHeader,
	diff: Difficulty,
	proof_size: usize,
	sz: u8,
) -> Result<(), Error> {
	let start_nonce = bh.pow.nonce;

	// set the nonce for faster solution finding in user testing
	if bh.height == 0 && global::is_user_testing_mode() {
		bh.pow.nonce = global::get_genesis_nonce();
	}

	// try to find a cuckoo cycle on that header hash
	loop {
		// if we found a cycle (not guaranteed) and the proof hash is higher that the
		// diff, we're all good
		let mut ctx = global::create_pow_context::<u32>(bh.height, sz, proof_size, MAX_SOLS)?;
		ctx.set_header_nonce(bh.pre_pow(), None, true)?;
		if let Ok(proofs) = ctx.find_cycles() {
			bh.pow.proof = proofs[0].clone();
			if bh.pow.to_difficulty(bh.height) >= diff {
				return Ok(());
			}
		}

		// otherwise increment the nonce
		let (res, _) = bh.pow.nonce.overflowing_add(1);
		bh.pow.nonce = res;

		// and if we're back where we started, update the time (changes the hash as
		// well)
		if bh.pow.nonce == start_nonce {
			bh.timestamp = DateTime::<Utc>::from_utc(NaiveDateTime::from_timestamp(0, 0), Utc);
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use crate::genesis;
	use crate::global;
	use crate::global::ChainTypes;

	/// We'll be generating genesis blocks differently
	#[test]
	fn genesis_pow() {
		global::set_mining_mode(ChainTypes::UserTesting);

		let mut b = genesis::genesis_dev();
		b.header.pow.nonce = 28106;
		b.header.pow.proof.edge_bits = global::min_edge_bits();
		println!("proof {}", global::proofsize());
		pow_size(
			&mut b.header,
			Difficulty::min(),
			global::proofsize(),
			global::min_edge_bits(),
		)
		.unwrap();
		println!("nonce {}", b.header.pow.nonce);
		assert_ne!(b.header.pow.nonce, 310);
		assert!(b.header.pow.to_difficulty(0) >= Difficulty::min());
		assert!(verify_size(&b.header).is_ok());
	}
}
<|endoftext|>


<commit_before>//! Unification and canonicalization logic.

use std::{borrow::Cow, fmt, sync::Arc};

use chalk_ir::{
    cast::Cast, fold::Fold, interner::HasInterner, FloatTy, IntTy, TyVariableKind, UniverseIndex,
    VariableKind,
};
use chalk_solve::infer::ParameterEnaVariableExt;
use ena::unify::UnifyKey;

use super::{InferOk, InferResult, InferenceContext, TypeError};
use crate::{
    db::HirDatabase, fold_tys, static_lifetime, BoundVar, Canonical, DebruijnIndex, GenericArg,
    InferenceVar, Interner, Scalar, Substitution, TraitEnvironment, Ty, TyKind,
};

impl<'a> InferenceContext<'a> {
    pub(super) fn canonicalize<T: Fold<Interner> + HasInterner<Interner = Interner>>(
        &mut self,
        t: T,
    ) -> Canonicalized<T::Result>
    where
        T::Result: HasInterner<Interner = Interner>,
    {
        let result = self.table.var_unification_table.canonicalize(&Interner, t);
        let free_vars = result
            .free_vars
            .into_iter()
            .map(|free_var| free_var.to_generic_arg(&Interner))
            .collect();
        Canonicalized { value: result.quantified, free_vars }
    }
}

#[derive(Debug)]
pub(super) struct Canonicalized<T>
where
    T: HasInterner<Interner = Interner>,
{
    pub(super) value: Canonical<T>,
    free_vars: Vec<GenericArg>,
}

impl<T: HasInterner<Interner = Interner>> Canonicalized<T> {
    pub(super) fn decanonicalize_ty(&self, ty: Ty) -> Ty {
        crate::fold_free_vars(ty, |bound, _binders| {
            let var = self.free_vars[bound.index].clone();
            var.assert_ty_ref(&Interner).clone()
        })
    }

    pub(super) fn apply_solution(
        &self,
        ctx: &mut InferenceContext<'_>,
        solution: Canonical<Substitution>,
    ) {
        // the solution may contain new variables, which we need to convert to new inference vars
        let new_vars = Substitution::from_iter(
            &Interner,
            solution.binders.iter(&Interner).map(|k| match k.kind {
                VariableKind::Ty(TyVariableKind::General) => {
                    ctx.table.new_type_var().cast(&Interner)
                }
                VariableKind::Ty(TyVariableKind::Integer) => {
                    ctx.table.new_integer_var().cast(&Interner)
                }
                VariableKind::Ty(TyVariableKind::Float) => {
                    ctx.table.new_float_var().cast(&Interner)
                }
                // Chalk can sometimes return new lifetime variables. We just use the static lifetime everywhere
                VariableKind::Lifetime => static_lifetime().cast(&Interner),
                _ => panic!("const variable in solution"),
            }),
        );
        for (i, ty) in solution.value.iter(&Interner).enumerate() {
            // FIXME: deal with non-type vars here -- the only problematic part is the normalization
            // and maybe we don't need that with lazy normalization?
            let var = self.free_vars[i].clone();
            // eagerly replace projections in the type; we may be getting types
            // e.g. from where clauses where this hasn't happened yet
            let ty = ctx.normalize_associated_types_in(
                new_vars.apply(ty.assert_ty_ref(&Interner).clone(), &Interner),
            );
            ctx.table.unify(var.assert_ty_ref(&Interner), &ty);
        }
    }
}

pub fn could_unify(db: &dyn HirDatabase, env: Arc<TraitEnvironment>, t1: &Ty, t2: &Ty) -> bool {
    InferenceTable::new(db, env).unify(t1, t2)
}

pub(crate) fn unify(
    db: &dyn HirDatabase,
    env: Arc<TraitEnvironment>,
    tys: &Canonical<(Ty, Ty)>,
) -> Option<Substitution> {
    let mut table = InferenceTable::new(db, env);
    let vars = Substitution::from_iter(
        &Interner,
        tys.binders
            .iter(&Interner)
            // we always use type vars here because we want everything to
            // fallback to Unknown in the end (kind of hacky, as below)
            .map(|_| table.new_type_var()),
    );
    let ty1_with_vars = vars.apply(tys.value.0.clone(), &Interner);
    let ty2_with_vars = vars.apply(tys.value.1.clone(), &Interner);
    if !table.unify(&ty1_with_vars, &ty2_with_vars) {
        return None;
    }
    // default any type vars that weren't unified back to their original bound vars
    // (kind of hacky)
    for (i, var) in vars.iter(&Interner).enumerate() {
        let var = var.assert_ty_ref(&Interner);
        if &*table.resolve_ty_shallow(var) == var {
            table.unify(
                var,
                &TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, i)).intern(&Interner),
            );
        }
    }
    Some(Substitution::from_iter(
        &Interner,
        vars.iter(&Interner)
            .map(|v| table.resolve_ty_completely(v.assert_ty_ref(&Interner).clone())),
    ))
}

#[derive(Clone, Debug)]
pub(super) struct TypeVariableTable {
    inner: Vec<TypeVariableData>,
}

impl TypeVariableTable {
    pub(super) fn set_diverging(&mut self, iv: InferenceVar, diverging: bool) {
        self.inner[iv.index() as usize].diverging = diverging;
    }

    fn fallback_value(&self, iv: InferenceVar, kind: TyVariableKind) -> Ty {
        match kind {
            _ if self.inner[iv.index() as usize].diverging => TyKind::Never,
            TyVariableKind::General => TyKind::Error,
            TyVariableKind::Integer => TyKind::Scalar(Scalar::Int(IntTy::I32)),
            TyVariableKind::Float => TyKind::Scalar(Scalar::Float(FloatTy::F64)),
        }
        .intern(&Interner)
    }
}

#[derive(Copy, Clone, Debug)]
pub(crate) struct TypeVariableData {
    diverging: bool,
}

type ChalkInferenceTable = chalk_solve::infer::InferenceTable<Interner>;

#[derive(Clone)]
pub(crate) struct InferenceTable<'a> {
    db: &'a dyn HirDatabase,
    trait_env: Arc<TraitEnvironment>,
    pub(super) var_unification_table: ChalkInferenceTable,
    pub(super) type_variable_table: TypeVariableTable,
}

impl<'a> InferenceTable<'a> {
    pub(crate) fn new(db: &'a dyn HirDatabase, trait_env: Arc<TraitEnvironment>) -> Self {
        InferenceTable {
            db,
            trait_env,
            var_unification_table: ChalkInferenceTable::new(),
            type_variable_table: TypeVariableTable { inner: Vec::new() },
        }
    }

    fn new_var(&mut self, kind: TyVariableKind, diverging: bool) -> Ty {
        let var = self.var_unification_table.new_variable(UniverseIndex::ROOT);
        self.type_variable_table.inner.extend(
            (0..1 + var.index() as usize - self.type_variable_table.inner.len())
                .map(|_| TypeVariableData { diverging: false }),
        );
        assert_eq!(var.index() as usize, self.type_variable_table.inner.len() - 1);
        self.type_variable_table.inner[var.index() as usize].diverging = diverging;
        var.to_ty_with_kind(&Interner, kind)
    }

    pub(crate) fn new_type_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::General, false)
    }

    pub(crate) fn new_integer_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::Integer, false)
    }

    pub(crate) fn new_float_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::Float, false)
    }

    pub(crate) fn new_maybe_never_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::General, true)
    }

    pub(crate) fn resolve_ty_completely(&mut self, ty: Ty) -> Ty {
        self.resolve_ty_completely_inner(&mut Vec::new(), ty)
    }

    // FIXME get rid of this, instead resolve shallowly where necessary
    pub(crate) fn resolve_ty_as_possible(&mut self, ty: Ty) -> Ty {
        self.resolve_ty_as_possible_inner(&mut Vec::new(), ty)
    }

    /// Unify two types and register new trait goals that arise from that.
    // TODO give these two functions better names
    pub(crate) fn unify(&mut self, ty1: &Ty, ty2: &Ty) -> bool {
        let _result = if let Ok(r) = self.unify_inner(ty1, ty2) {
            r
        } else {
            return false;
        };
        // TODO deal with new goals
        true
    }

    /// Unify two types and return new trait goals arising from it, so the
    /// caller needs to deal with them.
    pub(crate) fn unify_inner(&mut self, ty1: &Ty, ty2: &Ty) -> InferResult {
        match self.var_unification_table.relate(
            &Interner,
            &self.db,
            &self.trait_env.env,
            chalk_ir::Variance::Invariant,
            ty1,
            ty2,
        ) {
            Ok(_result) => {
                // TODO deal with new goals
                Ok(InferOk {})
            }
            Err(chalk_ir::NoSolution) => Err(TypeError),
        }
    }

    /// If `ty` is a type variable with known type, returns that type;
    /// otherwise, return ty.
    // FIXME this could probably just return Ty
    pub(crate) fn resolve_ty_shallow<'b>(&mut self, ty: &'b Ty) -> Cow<'b, Ty> {
        self.var_unification_table
            .normalize_ty_shallow(&Interner, ty)
            .map_or(Cow::Borrowed(ty), Cow::Owned)
    }

    /// Resolves the type as far as currently possible, replacing type variables
    /// by their known types. All types returned by the infer_* functions should
    /// be resolved as far as possible, i.e. contain no type variables with
    /// known type.
    fn resolve_ty_as_possible_inner(&mut self, tv_stack: &mut Vec<InferenceVar>, ty: Ty) -> Ty {
        fold_tys(
            ty,
            |ty, _| match ty.kind(&Interner) {
                &TyKind::InferenceVar(tv, kind) => {
                    if tv_stack.contains(&tv) {
                        cov_mark::hit!(type_var_cycles_resolve_as_possible);
                        // recursive type
                        return self.type_variable_table.fallback_value(tv, kind);
                    }
                    if let Some(known_ty) = self.var_unification_table.probe_var(tv) {
                        // known_ty may contain other variables that are known by now
                        tv_stack.push(tv);
                        let result = self.resolve_ty_as_possible_inner(
                            tv_stack,
                            known_ty.assert_ty_ref(&Interner).clone(),
                        );
                        tv_stack.pop();
                        result
                    } else {
                        ty
                    }
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }

    /// Resolves the type completely; type variables without known type are
    /// replaced by TyKind::Unknown.
    fn resolve_ty_completely_inner(&mut self, tv_stack: &mut Vec<InferenceVar>, ty: Ty) -> Ty {
        // FIXME implement as a proper Folder, handle lifetimes and consts as well
        fold_tys(
            ty,
            |ty, _| match ty.kind(&Interner) {
                &TyKind::InferenceVar(tv, kind) => {
                    if tv_stack.contains(&tv) {
                        cov_mark::hit!(type_var_cycles_resolve_completely);
                        // recursive type
                        return self.type_variable_table.fallback_value(tv, kind);
                    }
                    if let Some(known_ty) = self.var_unification_table.probe_var(tv) {
                        // known_ty may contain other variables that are known by now
                        tv_stack.push(tv);
                        let result = self.resolve_ty_completely_inner(
                            tv_stack,
                            known_ty.assert_ty_ref(&Interner).clone(),
                        );
                        tv_stack.pop();
                        result
                    } else {
                        self.type_variable_table.fallback_value(tv, kind)
                    }
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }
}

impl<'a> fmt::Debug for InferenceTable<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("InferenceTable")
            .field("num_vars", &self.type_variable_table.inner.len())
            .finish()
    }
}
<commit_msg>Fix fallback to bound vars in `unify`<commit_after>//! Unification and canonicalization logic.

use std::{borrow::Cow, fmt, sync::Arc};

use chalk_ir::{
    cast::Cast, fold::Fold, interner::HasInterner, FloatTy, IntTy, TyVariableKind, UniverseIndex,
};
use chalk_solve::infer::ParameterEnaVariableExt;
use ena::unify::UnifyKey;

use super::{InferOk, InferResult, InferenceContext, TypeError};
use crate::{
    db::HirDatabase, fold_tys, static_lifetime, BoundVar, Canonical, DebruijnIndex, GenericArg,
    InferenceVar, Interner, Scalar, Substitution, TraitEnvironment, Ty, TyKind, VariableKind,
};

impl<'a> InferenceContext<'a> {
    pub(super) fn canonicalize<T: Fold<Interner> + HasInterner<Interner = Interner>>(
        &mut self,
        t: T,
    ) -> Canonicalized<T::Result>
    where
        T::Result: HasInterner<Interner = Interner>,
    {
        let result = self.table.var_unification_table.canonicalize(&Interner, t);
        let free_vars = result
            .free_vars
            .into_iter()
            .map(|free_var| free_var.to_generic_arg(&Interner))
            .collect();
        Canonicalized { value: result.quantified, free_vars }
    }
}

#[derive(Debug)]
pub(super) struct Canonicalized<T>
where
    T: HasInterner<Interner = Interner>,
{
    pub(super) value: Canonical<T>,
    free_vars: Vec<GenericArg>,
}

impl<T: HasInterner<Interner = Interner>> Canonicalized<T> {
    pub(super) fn decanonicalize_ty(&self, ty: Ty) -> Ty {
        crate::fold_free_vars(ty, |bound, _binders| {
            let var = self.free_vars[bound.index].clone();
            var.assert_ty_ref(&Interner).clone()
        })
    }

    pub(super) fn apply_solution(
        &self,
        ctx: &mut InferenceContext<'_>,
        solution: Canonical<Substitution>,
    ) {
        // the solution may contain new variables, which we need to convert to new inference vars
        let new_vars = Substitution::from_iter(
            &Interner,
            solution.binders.iter(&Interner).map(|k| match k.kind {
                VariableKind::Ty(TyVariableKind::General) => {
                    ctx.table.new_type_var().cast(&Interner)
                }
                VariableKind::Ty(TyVariableKind::Integer) => {
                    ctx.table.new_integer_var().cast(&Interner)
                }
                VariableKind::Ty(TyVariableKind::Float) => {
                    ctx.table.new_float_var().cast(&Interner)
                }
                // Chalk can sometimes return new lifetime variables. We just use the static lifetime everywhere
                VariableKind::Lifetime => static_lifetime().cast(&Interner),
                _ => panic!("const variable in solution"),
            }),
        );
        for (i, ty) in solution.value.iter(&Interner).enumerate() {
            // FIXME: deal with non-type vars here -- the only problematic part is the normalization
            // and maybe we don't need that with lazy normalization?
            let var = self.free_vars[i].clone();
            // eagerly replace projections in the type; we may be getting types
            // e.g. from where clauses where this hasn't happened yet
            let ty = ctx.normalize_associated_types_in(
                new_vars.apply(ty.assert_ty_ref(&Interner).clone(), &Interner),
            );
            ctx.table.unify(var.assert_ty_ref(&Interner), &ty);
        }
    }
}

pub fn could_unify(db: &dyn HirDatabase, env: Arc<TraitEnvironment>, t1: &Ty, t2: &Ty) -> bool {
    InferenceTable::new(db, env).unify(t1, t2)
}

pub(crate) fn unify(
    db: &dyn HirDatabase,
    env: Arc<TraitEnvironment>,
    tys: &Canonical<(Ty, Ty)>,
) -> Option<Substitution> {
    let mut table = InferenceTable::new(db, env);
    let vars = Substitution::from_iter(
        &Interner,
        tys.binders
            .iter(&Interner)
            // we always use type vars here because we want everything to
            // fallback to Unknown in the end (kind of hacky, as below)
            .map(|_| table.new_type_var()),
    );
    let ty1_with_vars = vars.apply(tys.value.0.clone(), &Interner);
    let ty2_with_vars = vars.apply(tys.value.1.clone(), &Interner);
    if !table.unify(&ty1_with_vars, &ty2_with_vars) {
        return None;
    }
    // default any type vars that weren't unified back to their original bound vars
    // (kind of hacky)
    let find_var = |iv| {
        vars.iter(&Interner).position(|v| match v.interned() {
            chalk_ir::GenericArgData::Ty(ty) => ty.inference_var(&Interner),
            chalk_ir::GenericArgData::Lifetime(lt) => lt.inference_var(&Interner),
            chalk_ir::GenericArgData::Const(c) => c.inference_var(&Interner),
        } == Some(iv))
    };
    let fallback = |iv, kind, default| match kind {
        chalk_ir::VariableKind::Ty(_ty_kind) => find_var(iv).map_or(default, |i| {
            BoundVar::new(DebruijnIndex::INNERMOST, i).to_ty(&Interner).cast(&Interner)
        }),
        chalk_ir::VariableKind::Lifetime => find_var(iv).map_or(default, |i| {
            BoundVar::new(DebruijnIndex::INNERMOST, i).to_lifetime(&Interner).cast(&Interner)
        }),
        chalk_ir::VariableKind::Const(ty) => find_var(iv).map_or(default, |i| {
            BoundVar::new(DebruijnIndex::INNERMOST, i).to_const(&Interner, ty).cast(&Interner)
        }),
    };
    Some(Substitution::from_iter(
        &Interner,
        vars.iter(&Interner)
            .map(|v| table.resolve_with_fallback(v.assert_ty_ref(&Interner).clone(), fallback)),
    ))
}

#[derive(Clone, Debug)]
pub(super) struct TypeVariableTable {
    inner: Vec<TypeVariableData>,
}

impl TypeVariableTable {
    pub(super) fn set_diverging(&mut self, iv: InferenceVar, diverging: bool) {
        self.inner[iv.index() as usize].diverging = diverging;
    }

    fn fallback_value(&self, iv: InferenceVar, kind: TyVariableKind) -> Ty {
        match kind {
            _ if self.inner[iv.index() as usize].diverging => TyKind::Never,
            TyVariableKind::General => TyKind::Error,
            TyVariableKind::Integer => TyKind::Scalar(Scalar::Int(IntTy::I32)),
            TyVariableKind::Float => TyKind::Scalar(Scalar::Float(FloatTy::F64)),
        }
        .intern(&Interner)
    }
}

#[derive(Copy, Clone, Debug)]
pub(crate) struct TypeVariableData {
    diverging: bool,
}

type ChalkInferenceTable = chalk_solve::infer::InferenceTable<Interner>;

#[derive(Clone)]
pub(crate) struct InferenceTable<'a> {
    db: &'a dyn HirDatabase,
    trait_env: Arc<TraitEnvironment>,
    pub(super) var_unification_table: ChalkInferenceTable,
    pub(super) type_variable_table: TypeVariableTable,
}

impl<'a> InferenceTable<'a> {
    pub(crate) fn new(db: &'a dyn HirDatabase, trait_env: Arc<TraitEnvironment>) -> Self {
        InferenceTable {
            db,
            trait_env,
            var_unification_table: ChalkInferenceTable::new(),
            type_variable_table: TypeVariableTable { inner: Vec::new() },
        }
    }

    fn new_var(&mut self, kind: TyVariableKind, diverging: bool) -> Ty {
        let var = self.var_unification_table.new_variable(UniverseIndex::ROOT);
        self.type_variable_table.inner.extend(
            (0..1 + var.index() as usize - self.type_variable_table.inner.len())
                .map(|_| TypeVariableData { diverging: false }),
        );
        assert_eq!(var.index() as usize, self.type_variable_table.inner.len() - 1);
        self.type_variable_table.inner[var.index() as usize].diverging = diverging;
        var.to_ty_with_kind(&Interner, kind)
    }

    pub(crate) fn new_type_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::General, false)
    }

    pub(crate) fn new_integer_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::Integer, false)
    }

    pub(crate) fn new_float_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::Float, false)
    }

    pub(crate) fn new_maybe_never_var(&mut self) -> Ty {
        self.new_var(TyVariableKind::General, true)
    }

    pub(crate) fn resolve_with_fallback<T>(
        &mut self,
        t: T,
        fallback: impl Fn(InferenceVar, VariableKind, GenericArg) -> GenericArg,
    ) -> T::Result
    where
        T: HasInterner<Interner = Interner> + Fold<Interner>,
    {
        self.resolve_with_fallback_inner(&mut Vec::new(), t, &fallback)
    }

    fn resolve_with_fallback_inner<T>(
        &mut self,
        var_stack: &mut Vec<InferenceVar>,
        t: T,
        fallback: &impl Fn(InferenceVar, VariableKind, GenericArg) -> GenericArg,
    ) -> T::Result
    where
        T: HasInterner<Interner = Interner> + Fold<Interner>,
    {
        fold_tys(
            t,
            |ty, _| match ty.kind(&Interner) {
                &TyKind::InferenceVar(tv, kind) => {
                    if var_stack.contains(&tv) {
                        cov_mark::hit!(type_var_cycles_resolve_as_possible);
                        // recursive type
                        let default =
                            self.type_variable_table.fallback_value(tv, kind).cast(&Interner);
                        return fallback(tv, VariableKind::Ty(kind), default)
                            .assert_ty_ref(&Interner)
                            .clone();
                    }
                    if let Some(known_ty) = self.var_unification_table.probe_var(tv) {
                        // known_ty may contain other variables that are known by now
                        var_stack.push(tv);
                        let result = self.resolve_with_fallback_inner(
                            var_stack,
                            known_ty.assert_ty_ref(&Interner).clone(),
                            fallback,
                        );
                        var_stack.pop();
                        result
                    } else {
                        let default =
                            self.type_variable_table.fallback_value(tv, kind).cast(&Interner);
                        fallback(tv, VariableKind::Ty(kind), default)
                            .assert_ty_ref(&Interner)
                            .clone()
                    }
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }

    pub(crate) fn resolve_ty_completely(&mut self, ty: Ty) -> Ty {
        self.resolve_with_fallback(ty, |_, _, d| d)
    }

    // FIXME get rid of this, instead resolve shallowly where necessary
    pub(crate) fn resolve_ty_as_possible(&mut self, ty: Ty) -> Ty {
        self.resolve_ty_as_possible_inner(&mut Vec::new(), ty)
    }

    /// Unify two types and register new trait goals that arise from that.
    // TODO give these two functions better names
    pub(crate) fn unify(&mut self, ty1: &Ty, ty2: &Ty) -> bool {
        let _result = if let Ok(r) = self.unify_inner(ty1, ty2) {
            r
        } else {
            return false;
        };
        // TODO deal with new goals
        true
    }

    /// Unify two types and return new trait goals arising from it, so the
    /// caller needs to deal with them.
    pub(crate) fn unify_inner(&mut self, ty1: &Ty, ty2: &Ty) -> InferResult {
        match self.var_unification_table.relate(
            &Interner,
            &self.db,
            &self.trait_env.env,
            chalk_ir::Variance::Invariant,
            ty1,
            ty2,
        ) {
            Ok(_result) => {
                // TODO deal with new goals
                Ok(InferOk {})
            }
            Err(chalk_ir::NoSolution) => Err(TypeError),
        }
    }

    /// If `ty` is a type variable with known type, returns that type;
    /// otherwise, return ty.
    // FIXME this could probably just return Ty
    pub(crate) fn resolve_ty_shallow<'b>(&mut self, ty: &'b Ty) -> Cow<'b, Ty> {
        self.var_unification_table
            .normalize_ty_shallow(&Interner, ty)
            .map_or(Cow::Borrowed(ty), Cow::Owned)
    }

    /// Resolves the type as far as currently possible, replacing type variables
    /// by their known types. All types returned by the infer_* functions should
    /// be resolved as far as possible, i.e. contain no type variables with
    /// known type.
    fn resolve_ty_as_possible_inner(&mut self, tv_stack: &mut Vec<InferenceVar>, ty: Ty) -> Ty {
        fold_tys(
            ty,
            |ty, _| match ty.kind(&Interner) {
                &TyKind::InferenceVar(tv, kind) => {
                    if tv_stack.contains(&tv) {
                        cov_mark::hit!(type_var_cycles_resolve_as_possible);
                        // recursive type
                        return self.type_variable_table.fallback_value(tv, kind);
                    }
                    if let Some(known_ty) = self.var_unification_table.probe_var(tv) {
                        // known_ty may contain other variables that are known by now
                        tv_stack.push(tv);
                        let result = self.resolve_ty_as_possible_inner(
                            tv_stack,
                            known_ty.assert_ty_ref(&Interner).clone(),
                        );
                        tv_stack.pop();
                        result
                    } else {
                        ty
                    }
                }
                _ => ty,
            },
            DebruijnIndex::INNERMOST,
        )
    }
}

impl<'a> fmt::Debug for InferenceTable<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("InferenceTable")
            .field("num_vars", &self.type_variable_table.inner.len())
            .finish()
    }
}
<|endoftext|>


<commit_before>// lock-free stack
use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::SeqCst,
};

use super::*;

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub struct Node<T: Send + 'static> {
    pub(crate) inner: T,
    pub(crate) next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let next = self.next.load(SeqCst, unprotected()).as_raw();
            if !next.is_null() {
                drop(Box::from_raw(next as *mut Node<T>));
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Stack<T> {
        Stack {
            head: Atomic::null(),
        }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr = self.head.load(SeqCst, unprotected()).as_raw();
            if !curr.is_null() {
                drop(Box::from_raw(curr as *mut Node<T>));
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Clone + Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let guard = pin();
        let head = self.head(&guard);
        let iter = StackIter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter.write_str(&*format!("({:?}) ", &node as *const _))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Clone + Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub fn push(&self, inner: T) {
        debug_delay();
        let node = Owned::new(Node {
            inner,
            next: Atomic::null(),
        });

        unsafe {
            let node = node.into_shared(unprotected());

            loop {
                let head = self.head(unprotected());
                node.deref().next.store(head, SeqCst);
                if self
                    .head
                    .compare_and_set(head, node, SeqCst, unprotected())
                    .is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    fn _pop<'g>(&self, guard: &'g Guard) -> Option<T> {
        use std::ptr;
        debug_delay();
        let mut head = self.head(&guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(SeqCst, &guard);
                    match self.head.compare_and_set(head, next, SeqCst, &guard)
                    {
                        Ok(_) => unsafe {
                            guard.defer_destroy(head);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// compare and push
    pub fn cap<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        new: T,
        guard: &'g Guard,
    ) -> std::result::Result<
        Shared<'g, Node<T>>,
        (Shared<'g, Node<T>>, Owned<Node<T>>),
    > {
        debug_delay();
        let node = Owned::new(Node {
            inner: new,
            next: Atomic::from(old),
        });

        self.cap_node(old, node, guard)
    }

    /// compare and push
    pub fn cap_node<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        mut node: Owned<Node<T>>,
        guard: &'g Guard,
    ) -> std::result::Result<
        Shared<'g, Node<T>>,
        (Shared<'g, Node<T>>, Owned<Node<T>>),
    > {
        // properly set next ptr
        node.next = Atomic::from(old);
        let res = self.head.compare_and_set(old, node, SeqCst, guard);

        match res {
            Err(e) => {
                // we want to set next to null to prevent
                // the current shared head from being
                // dropped when we drop this node.
                let mut returned = e.new;
                returned.next = Atomic::null();
                Err((e.current, returned))
            }
            Ok(success) => Ok(success),
        }
    }

    /// compare and swap
    pub fn cas<'g>(
        &self,
        old: Shared<'g, Node<T>>,
        new: Owned<Node<T>>,
        guard: &'g Guard,
    ) -> std::result::Result<
        Shared<'g, Node<T>>,
        (Shared<'g, Node<T>>, Owned<Node<T>>),
    > {
        debug_delay();
        let res = self.head.compare_and_set(old, new, SeqCst, guard);

        match res {
            Ok(success) => {
                if !old.is_null() {
                    unsafe {
                        guard.defer_destroy(old);
                    };
                }
                Ok(success)
            }
            Err(e) => Err((e.current, e.new)),
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub fn head<'g>(&self, guard: &'g Guard) -> Shared<'g, Node<T>> {
        self.head.load(SeqCst, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub struct StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> StackIter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a StackIter from a pointer to one.
    pub fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> StackIter<'b, T> {
        StackIter { inner: ptr, guard }
    }
}

impl<'a, T> Iterator for StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner = self.inner.deref().next.load(SeqCst, self.guard);
                Some(ret)
            }
        }
    }
}

/// Turns a vector of elements into a lock-free stack
/// of them, and returns the head of the stack.
pub fn node_from_frag_vec<T>(from: Vec<T>) -> Owned<Node<T>>
where
    T: Send + 'static + Sync,
{
    let mut last = None;

    for item in from.into_iter().rev() {
        last = if let Some(last) = last {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::from(last),
            }))
        } else {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::null(),
            }))
        }
    }

    last.expect("at least one frag was provided in the from Vec")
}

#[test]
fn basic_functionality() {
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll._pop(&guard), None);
    ll.push(1);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        ll2.push(2);
        ll2.push(3);
        ll2.push(4);
    });
    t.join().unwrap();
    ll.push(5);
    assert_eq!(ll._pop(&guard), Some(5));
    assert_eq!(ll._pop(&guard), Some(4));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3._pop(&guard), Some(3));
        assert_eq!(ll3._pop(&guard), Some(2));
    });
    t.join().unwrap();
    assert_eq!(ll._pop(&guard), Some(1));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4._pop(&guard), None);
    });
    t.join().unwrap();
}
<commit_msg>Extract complex return type for Stack atomic operations<commit_after>// lock-free stack
use std::{
    fmt::{self, Debug},
    ops::Deref,
    sync::atomic::Ordering::SeqCst,
};

use super::*;

type CompareAndSwapResult<'g, T> = std::result::Result<
    Shared<'g, Node<T>>,
    (Shared<'g, Node<T>>, Owned<Node<T>>),
>;

/// A node in the lock-free `Stack`.
#[derive(Debug)]
pub struct Node<T: Send + 'static> {
    pub(crate) inner: T,
    pub(crate) next: Atomic<Node<T>>,
}

impl<T: Send + 'static> Drop for Node<T> {
    fn drop(&mut self) {
        unsafe {
            let next = self.next.load(SeqCst, unprotected()).as_raw();
            if !next.is_null() {
                drop(Box::from_raw(next as *mut Node<T>));
            }
        }
    }
}

/// A simple lock-free stack, with the ability to atomically
/// append or entirely swap-out entries.
pub struct Stack<T: Send + 'static> {
    head: Atomic<Node<T>>,
}

impl<T: Send + 'static> Default for Stack<T> {
    fn default() -> Stack<T> {
        Stack {
            head: Atomic::null(),
        }
    }
}

impl<T: Send + 'static> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            let curr = self.head.load(SeqCst, unprotected()).as_raw();
            if !curr.is_null() {
                drop(Box::from_raw(curr as *mut Node<T>));
            }
        }
    }
}

impl<T> Debug for Stack<T>
where
    T: Clone + Debug + Send + 'static + Sync,
{
    fn fmt(
        &self,
        formatter: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let guard = pin();
        let head = self.head(&guard);
        let iter = StackIter::from_ptr(head, &guard);

        formatter.write_str("Stack [")?;
        let mut written = false;
        for node in iter {
            if written {
                formatter.write_str(", ")?;
            }
            formatter.write_str(&*format!("({:?}) ", &node as *const _))?;
            node.fmt(formatter)?;
            written = true;
        }
        formatter.write_str("]")?;
        Ok(())
    }
}

impl<T: Send + 'static> Deref for Node<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.inner
    }
}

impl<T: Clone + Send + Sync + 'static> Stack<T> {
    /// Add an item to the stack, spinning until successful.
    pub fn push(&self, inner: T) {
        debug_delay();
        let node = Owned::new(Node {
            inner,
            next: Atomic::null(),
        });

        unsafe {
            let node = node.into_shared(unprotected());

            loop {
                let head = self.head(unprotected());
                node.deref().next.store(head, SeqCst);
                if self
                    .head
                    .compare_and_set(head, node, SeqCst, unprotected())
                    .is_ok()
                {
                    return;
                }
            }
        }
    }

    /// Pop the next item off the stack. Returns None if nothing is there.
    fn _pop<'g>(&self, guard: &'g Guard) -> Option<T> {
        use std::ptr;
        debug_delay();
        let mut head = self.head(&guard);
        loop {
            match unsafe { head.as_ref() } {
                Some(h) => {
                    let next = h.next.load(SeqCst, &guard);
                    match self.head.compare_and_set(head, next, SeqCst, &guard)
                    {
                        Ok(_) => unsafe {
                            guard.defer_destroy(head);
                            return Some(ptr::read(&h.inner));
                        },
                        Err(h) => head = h.current,
                    }
                }
                None => return None,
            }
        }
    }

    /// compare and push
    pub fn cap<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        new: T,
        guard: &'g Guard,
    ) -> CompareAndSwapResult<'g, T> {
        debug_delay();
        let node = Owned::new(Node {
            inner: new,
            next: Atomic::from(old),
        });

        self.cap_node(old, node, guard)
    }

    /// compare and push
    pub fn cap_node<'g>(
        &self,
        old: Shared<'_, Node<T>>,
        mut node: Owned<Node<T>>,
        guard: &'g Guard,
    ) -> CompareAndSwapResult<'g, T> {
        // properly set next ptr
        node.next = Atomic::from(old);
        let res = self.head.compare_and_set(old, node, SeqCst, guard);

        match res {
            Err(e) => {
                // we want to set next to null to prevent
                // the current shared head from being
                // dropped when we drop this node.
                let mut returned = e.new;
                returned.next = Atomic::null();
                Err((e.current, returned))
            }
            Ok(success) => Ok(success),
        }
    }

    /// compare and swap
    pub fn cas<'g>(
        &self,
        old: Shared<'g, Node<T>>,
        new: Owned<Node<T>>,
        guard: &'g Guard,
    ) -> CompareAndSwapResult<'g, T> {
        debug_delay();
        let res = self.head.compare_and_set(old, new, SeqCst, guard);

        match res {
            Ok(success) => {
                if !old.is_null() {
                    unsafe {
                        guard.defer_destroy(old);
                    };
                }
                Ok(success)
            }
            Err(e) => Err((e.current, e.new)),
        }
    }

    /// Returns the current head pointer of the stack, which can
    /// later be used as the key for cas and cap operations.
    pub fn head<'g>(&self, guard: &'g Guard) -> Shared<'g, Node<T>> {
        self.head.load(SeqCst, guard)
    }
}

/// An iterator over nodes in a lock-free stack.
pub struct StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    inner: Shared<'a, Node<T>>,
    guard: &'a Guard,
}

impl<'a, T> StackIter<'a, T>
where
    T: 'a + Send + 'static + Sync,
{
    /// Creates a StackIter from a pointer to one.
    pub fn from_ptr<'b>(
        ptr: Shared<'b, Node<T>>,
        guard: &'b Guard,
    ) -> StackIter<'b, T> {
        StackIter { inner: ptr, guard }
    }
}

impl<'a, T> Iterator for StackIter<'a, T>
where
    T: Send + 'static + Sync,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        debug_delay();
        if self.inner.is_null() {
            None
        } else {
            unsafe {
                let ret = &self.inner.deref().inner;
                self.inner = self.inner.deref().next.load(SeqCst, self.guard);
                Some(ret)
            }
        }
    }
}

/// Turns a vector of elements into a lock-free stack
/// of them, and returns the head of the stack.
pub fn node_from_frag_vec<T>(from: Vec<T>) -> Owned<Node<T>>
where
    T: Send + 'static + Sync,
{
    let mut last = None;

    for item in from.into_iter().rev() {
        last = if let Some(last) = last {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::from(last),
            }))
        } else {
            Some(Owned::new(Node {
                inner: item,
                next: Atomic::null(),
            }))
        }
    }

    last.expect("at least one frag was provided in the from Vec")
}

#[test]
fn basic_functionality() {
    use std::sync::Arc;
    use std::thread;

    let guard = pin();
    let ll = Arc::new(Stack::default());
    assert_eq!(ll._pop(&guard), None);
    ll.push(1);
    let ll2 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        ll2.push(2);
        ll2.push(3);
        ll2.push(4);
    });
    t.join().unwrap();
    ll.push(5);
    assert_eq!(ll._pop(&guard), Some(5));
    assert_eq!(ll._pop(&guard), Some(4));
    let ll3 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll3._pop(&guard), Some(3));
        assert_eq!(ll3._pop(&guard), Some(2));
    });
    t.join().unwrap();
    assert_eq!(ll._pop(&guard), Some(1));
    let ll4 = Arc::clone(&ll);
    let t = thread::spawn(move || {
        let guard = pin();
        assert_eq!(ll4._pop(&guard), None);
    });
    t.join().unwrap();
}
<|endoftext|>