mentatdb 1.18.0

// Copyright 2018 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.

#[macro_use]
extern crate mentat;

use std::collections::{BTreeMap, BTreeSet};

use std::path::{Path, PathBuf};

use core_traits::{Binding, StructuredMap};

// TODO: re-export from Mentat.
use mentat_core::ValueRc;

use mentat::{
    Entid, HasSchema, IntoResult, Keyword, Pullable, QueryInputs, Queryable, RelResult, Store,
    TypedValue,
};

fn fixture_path(rest: &str) -> PathBuf {
    let fixtures = Path::new("../fixtures/");
    fixtures.join(Path::new(rest))
}

#[test]
fn test_simple_pull() {
    let beacon;
    let capitol;
    let beacon_district;
    let capitol_district;

    let mut store = Store::open("").expect("opened");
    {
        let mut in_progress = store.begin_transaction().expect("began");
        in_progress
            .import(fixture_path("cities.schema"))
            .expect("transacted schema");
        let report = in_progress
            .import(fixture_path("all_seattle.edn"))
            .expect("transacted data");

        // These tempid strings come out of all_seattle.edn.
        beacon = *report
            .tempids
            .get(&"a17592186045471".to_string())
            .expect("beacon");
        beacon_district = *report
            .tempids
            .get(&"a17592186045450".to_string())
            .expect("beacon district");

        capitol = *report
            .tempids
            .get(&"a17592186045439".to_string())
            .expect("capitol");
        capitol_district = *report
            .tempids
            .get(&"a17592186045438".to_string())
            .expect("capitol district");

        in_progress.commit().expect("committed");
    }

    let schema = store.conn().current_schema();
    let district = schema
        .get_entid(&kw!(:neighborhood/district))
        .expect("district")
        .0;
    let name = schema.get_entid(&kw!(:neighborhood/name)).expect("name").0;
    let attrs: BTreeSet<Entid> = vec![district, name].into_iter().collect();

    // Find Beacon Hill and Capitol Hill via query.
    let query = r#"[:find [?hood ...]
                    :where
                    (or [?hood :neighborhood/name "Beacon Hill"]
                        [?hood :neighborhood/name "Capitol Hill"])]"#;
    let hoods: BTreeSet<Entid> = store
        .q_once(query, None)
        .into_coll_result()
        .expect("hoods")
        .into_iter()
        .map(|b| {
            b.into_scalar()
                .and_then(|tv| tv.into_entid())
                .expect("scalar")
        })
        .collect();

    println!("Hoods: {:?}", hoods);

    // The query and the tempids agree.
    assert_eq!(hoods, vec![beacon, capitol].into_iter().collect());

    // Pull attributes of those two neighborhoods.
    let pulled = store
        .begin_read()
        .expect("read")
        .pull_attributes_for_entities(hoods, attrs)
        .expect("pulled");

    // Here's what we expect:
    let c: StructuredMap = vec![
        (kw!(:neighborhood/name), "Capitol Hill".into()),
        (
            kw!(:neighborhood/district),
            TypedValue::Ref(capitol_district),
        ),
    ]
    .into();
    let b: StructuredMap = vec![
        (kw!(:neighborhood/name), "Beacon Hill".into()),
        (
            kw!(:neighborhood/district),
            TypedValue::Ref(beacon_district),
        ),
    ]
    .into();

    let expected: BTreeMap<Entid, ValueRc<StructuredMap>> =
        vec![(capitol, c.into()), (beacon, b.into())]
            .into_iter()
            .collect();

    assert_eq!(pulled, expected);

    // Now test pull inside the query itself.
    let query = r#"[:find ?hood (pull ?district [:db/id
                                                 :district/name :as :district/district
                                                 :district/region])
                    :where
                    (or [?hood :neighborhood/name "Beacon Hill"]
                        [?hood :neighborhood/name "Capitol Hill"])
                    [?hood :neighborhood/district ?district]
                    :order ?hood]"#;
    let results: RelResult<Binding> = store
        .begin_read()
        .expect("read")
        .q_once(query, None)
        .into_rel_result()
        .expect("results");

    let beacon_district_pull: Vec<(Keyword, TypedValue)> = vec![
        (kw!(:db/id), TypedValue::Ref(beacon_district)),
        (kw!(:district/district), "Greater Duwamish".into()),
        (
            kw!(:district/region),
            schema
                .get_entid(&Keyword::namespaced("region", "se"))
                .unwrap()
                .into(),
        ),
    ];
    let beacon_district_pull: StructuredMap = beacon_district_pull.into();
    let capitol_district_pull: Vec<(Keyword, TypedValue)> = vec![
        (kw!(:db/id), TypedValue::Ref(capitol_district)),
        (kw!(:district/district), "East".into()),
        (
            kw!(:district/region),
            schema
                .get_entid(&Keyword::namespaced("region", "e"))
                .unwrap()
                .into(),
        ),
    ];
    let capitol_district_pull: StructuredMap = capitol_district_pull.into();

    let expected = RelResult {
        width: 2,
        values: vec![
            TypedValue::Ref(capitol).into(),
            capitol_district_pull.into(),
            TypedValue::Ref(beacon).into(),
            beacon_district_pull.into(),
        ],
    };
    assert_eq!(results, expected);

    // We can also prepare the query.
    let reader = store.begin_read().expect("read");
    let mut prepared = reader.q_prepare(query, None).expect("prepared");

    assert_eq!(
        prepared.run(None).into_rel_result().expect("results"),
        expected
    );

    // Execute a scalar query where the body is constant.
    // TODO: we shouldn't require `:where`; that makes this non-constant!
    let query = r#"[:find (pull ?hood [:db/id :as :neighborhood
                                       :neighborhood/name]) .
                    :in ?hood
                    :where [?hood :neighborhood/district _]]"#;
    let result = reader
        .q_once(
            query,
            QueryInputs::with_value_sequence(vec![(var!(?hood), TypedValue::Ref(beacon))]),
        )
        .into_scalar_result()
        .expect("success")
        .expect("result");

    let expected: StructuredMap = vec![
        (kw!(:neighborhood/name), TypedValue::from("Beacon Hill")),
        (kw!(:neighborhood), TypedValue::Ref(beacon)),
    ]
    .into();
    assert_eq!(result, expected.into());

    // Collect the names and regions of all districts.
    let query = r#"[:find [(pull ?district [:district/name :district/region]) ...]
                    :where [_ :neighborhood/district ?district]]"#;
    let results = reader
        .q_once(query, None)
        .into_coll_result()
        .expect("result");

    let expected: Vec<StructuredMap> = vec![
        vec![
            (kw!(:district/name), TypedValue::from("East")),
            (kw!(:district/region), TypedValue::Ref(65556)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Southwest")),
            (kw!(:district/region), TypedValue::Ref(65559)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Downtown")),
            (kw!(:district/region), TypedValue::Ref(65560)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Greater Duwamish")),
            (kw!(:district/region), TypedValue::Ref(65557)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Ballard")),
            (kw!(:district/region), TypedValue::Ref(65561)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Northeast")),
            (kw!(:district/region), TypedValue::Ref(65555)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Southeast")),
            (kw!(:district/region), TypedValue::Ref(65557)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Northwest")),
            (kw!(:district/region), TypedValue::Ref(65561)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Central")),
            (kw!(:district/region), TypedValue::Ref(65556)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Delridge")),
            (kw!(:district/region), TypedValue::Ref(65559)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Lake Union")),
            (kw!(:district/region), TypedValue::Ref(65560)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("Magnolia/Queen Anne")),
            (kw!(:district/region), TypedValue::Ref(65560)),
        ]
        .into(),
        vec![
            (kw!(:district/name), TypedValue::from("North")),
            (kw!(:district/region), TypedValue::Ref(65555)),
        ]
        .into(),
    ];

    let expected: Vec<Binding> = expected.into_iter().map(|m| m.into()).collect();
    assert_eq!(results, expected);

    // Pull fulltext.
    let query = r#"[:find [(pull ?c [:community/name :community/category]) ...]
                    :where
                    [?c :community/name ?name]
                    [?c :community/type :community.type/website]
                    [(fulltext $ :community/category "food") [[?c ?cat]]]]"#;
    let results = reader
        .q_once(query, None)
        .into_coll_result()
        .expect("result");
    let expected: Vec<StructuredMap> = vec![
        vec![
            (
                kw!(:community/name),
                Binding::Scalar(TypedValue::from("Community Harvest of Southwest Seattle")),
            ),
            (
                kw!(:community/category),
                vec![Binding::Scalar(TypedValue::from("sustainable food"))].into(),
            ),
        ]
        .into(),
        vec![
            (
                kw!(:community/name),
                Binding::Scalar(TypedValue::from("InBallard")),
            ),
            (
                kw!(:community/category),
                vec![
                    Binding::Scalar(TypedValue::from("shopping")),
                    Binding::Scalar(TypedValue::from("food")),
                    Binding::Scalar(TypedValue::from("nightlife")),
                    Binding::Scalar(TypedValue::from("services")),
                ]
                .into(),
            ),
        ]
        .into(),
        vec![
            (
                kw!(:community/name),
                Binding::Scalar(TypedValue::from("Seattle Chinatown Guide")),
            ),
            (
                kw!(:community/category),
                vec![
                    Binding::Scalar(TypedValue::from("shopping")),
                    Binding::Scalar(TypedValue::from("food")),
                ]
                .into(),
            ),
        ]
        .into(),
        vec![
            (
                kw!(:community/name),
                Binding::Scalar(TypedValue::from("University District Food Bank")),
            ),
            (
                kw!(:community/category),
                vec![Binding::Scalar(TypedValue::from("food bank"))].into(),
            ),
        ]
        .into(),
    ];

    let expected: Vec<Binding> = expected.into_iter().map(|m| m.into()).collect();
    assert_eq!(results, expected);
}

// TEST:
// - Constant query bodies in pull.
// - Values that are present in the cache (=> constant pull, too).
// - Pull for each kind of find spec.
// - That the keys in each map are ValueRc::ptr_eq.
// - Entity presence/absence when the pull expressions don't match anything.
// - Aliases. (No parser support yet.)

// Ported from pg_mentat: reverse-reference pull (`:ns/_attr`).
// `(pull ?e [:person/_friend])` on entity ?e returns every entity that has ?e
// as a `:person/friend` value, as a cardinality-many list under `:person/_friend`.
#[test]
fn test_reverse_pull() {
    let mut store = Store::open("").expect("opened");

    let (alice, bob, carol) = {
        let mut in_progress = store.begin_transaction().expect("began");
        in_progress
            .transact(
                r#"[{:db/ident :person/name
                     :db/valueType :db.type/string
                     :db/cardinality :db.cardinality/one}
                    {:db/ident :person/friend
                     :db/valueType :db.type/ref
                     :db/cardinality :db.cardinality/many}]"#,
            )
            .expect("schema");

        // Bob and Carol both consider Alice a friend; Alice has no friends.
        let report = in_progress
            .transact(
                r#"[{:db/id "alice" :person/name "Alice"}
                    {:db/id "bob"   :person/name "Bob"   :person/friend "alice"}
                    {:db/id "carol" :person/name "Carol" :person/friend "alice"}]"#,
            )
            .expect("data");
        let a = *report.tempids.get("alice").expect("alice");
        let b = *report.tempids.get("bob").expect("bob");
        let c = *report.tempids.get("carol").expect("carol");
        in_progress.commit().expect("committed");
        (a, b, c)
    };

    let reader = store.begin_read().expect("read");

    // Reverse: who considers Alice a friend? => Bob and Carol (order by e asc).
    let rev = reader
        .q_once(
            r#"[:find (pull ?e [:person/_friend]) . :in ?e :where [?e :person/name _]]"#,
            QueryInputs::with_value_sequence(vec![(var!(?e), TypedValue::Ref(alice))]),
        )
        .into_scalar_result()
        .expect("ok")
        .expect("map");

    // Bob and Carol were transacted after Alice, so bob < carol by entid.
    let (lo, hi) = if bob < carol {
        (bob, carol)
    } else {
        (carol, bob)
    };
    let expected: StructuredMap = vec![(
        kw!(:person/_friend),
        Binding::Vec(ValueRc::new(vec![
            Binding::Scalar(TypedValue::Ref(lo)),
            Binding::Scalar(TypedValue::Ref(hi)),
        ])),
    )]
    .into();
    assert_eq!(rev, expected.into());

    // Bob refers to nobody, so a reverse pull on Bob yields an empty map.
    let none = reader
        .q_once(
            r#"[:find (pull ?e [:person/_friend]) . :in ?e :where [?e :person/name _]]"#,
            QueryInputs::with_value_sequence(vec![(var!(?e), TypedValue::Ref(bob))]),
        )
        .into_scalar_result()
        .expect("ok")
        .expect("map");
    let empty: StructuredMap = Vec::<(Keyword, Binding)>::new().into();
    assert_eq!(none, empty.into());
}

// Map specs `{:attr [sub-pattern]}`: follow a ref and pull the sub-pattern on
// each target. Each nested map carries `:db/id` (pg_mentat's edn_pull shape);
// cardinality-one nests a map, cardinality-many and reverse a vector of maps.
#[test]
fn test_nested_map_pull() {
    let mut store = Store::open("").expect("opened");
    store
        .transact(
            r#"[{:db/ident :person/name :db/valueType :db.type/string :db/cardinality :db.cardinality/one}
                {:db/ident :person/friend :db/valueType :db.type/ref :db/cardinality :db.cardinality/many}
                {:db/ident :person/pet :db/valueType :db.type/ref :db/cardinality :db.cardinality/one}
                {:db/ident :pet/name :db/valueType :db.type/string :db/cardinality :db.cardinality/one}]"#,
        )
        .expect("schema");
    let r = store
        .transact(
            r#"[{:db/id "rex" :pet/name "Rex"}
                {:db/id "alice" :person/name "Alice" :person/pet "rex"}
                {:db/id "bob" :person/name "Bob" :person/friend "alice"}
                {:db/id "carol" :person/name "Carol" :person/friend ["alice" "bob"]}]"#,
        )
        .expect("data");
    let id = |k: &str| r.tempids[k];
    let (alice, bob, carol, rex) = (id("alice"), id("bob"), id("carol"), id("rex"));
    let t1 = r.tx_id;

    let ref_ = |e: Entid| Binding::Scalar(TypedValue::Ref(e));
    let s = |v: &str| Binding::Scalar(v.into());
    let map = |kvs: Vec<(Keyword, Binding)>| -> Binding {
        let m: StructuredMap = kvs.into();
        m.into()
    };
    let db_id = kw!(:db/id);

    // Cardinality-many forward, with a cardinality-one map spec inside it.
    let carol_now = store
        .pull(
            carol,
            "[:person/name {:person/friend [:person/name {:person/pet [:pet/name]}]}]",
        )
        .expect("pull");
    let (lo, hi) = if alice < bob {
        (alice, bob)
    } else {
        (bob, alice)
    };
    let friend = |e: Entid| -> Binding {
        if e == alice {
            map(vec![
                (db_id.clone(), ref_(alice)),
                (kw!(:person/name), s("Alice")),
                (
                    kw!(:person/pet),
                    map(vec![(db_id.clone(), ref_(rex)), (kw!(:pet/name), s("Rex"))]),
                ),
            ])
        } else {
            map(vec![
                (db_id.clone(), ref_(bob)),
                (kw!(:person/name), s("Bob")),
            ])
        }
    };
    let expected: StructuredMap = vec![
        (kw!(:person/name), s("Carol")),
        (
            kw!(:person/friend),
            Binding::Vec(ValueRc::new(vec![friend(lo), friend(hi)])),
        ),
    ]
    .into();
    assert_eq!(carol_now, expected);

    // Reverse: who has Alice as a friend, with their names.
    let alice_rev = store
        .pull(alice, "[{:person/_friend [:person/name]}]")
        .expect("pull");
    let who = match alice_rev.get(&kw!(:person/_friend)) {
        Some(Binding::Vec(vs)) => vs
            .iter()
            .map(|b| match b {
                Binding::Map(m) => m.get(&kw!(:person/name)).cloned().unwrap(),
                other => panic!("{other:?}"),
            })
            .collect::<Vec<_>>(),
        other => panic!("expected a vec of maps, got {other:?}"),
    };
    assert_eq!(who.len(), 2);
    assert!(
        who.contains(&s("Bob")) && who.contains(&s("Carol")),
        "{who:?}"
    );

    // `*` with a map spec: the spec's nesting replaces the plain ref.
    let alice_star = store
        .pull(alice, "[* {:person/pet [:pet/name]}]")
        .expect("pull");
    assert_eq!(alice_star.get(&kw!(:person/name)), Some(&s("Alice")));
    assert_eq!(
        alice_star.get(&kw!(:person/pet)),
        Some(&map(vec![
            (db_id.clone(), ref_(rex)),
            (kw!(:pet/name), s("Rex"))
        ]))
    );

    // A map spec on an entity without the ref is absent, not empty.
    let rex_pull = store
        .pull(rex, "[:pet/name {:person/pet [:pet/name]}]")
        .expect("pull");
    assert_eq!(rex_pull.get(&kw!(:person/pet)), None);

    // As a find-pull, and on an as-of basis: Rex renamed after t1.
    store
        .transact(&format!("[[:db/add {rex} :pet/name \"Max\"]]"))
        .expect("rename");
    let q = r#"[:find (pull ?e [{:person/pet [:pet/name]}]) . :where [?e :person/name "Alice"]]"#;
    let pet_name = |out: Option<Binding>| match out {
        Some(Binding::Map(m)) => match m.get(&kw!(:person/pet)) {
            Some(Binding::Map(p)) => p.get(&kw!(:pet/name)).cloned(),
            other => panic!("{other:?}"),
        },
        other => panic!("{other:?}"),
    };
    assert_eq!(
        pet_name(store.q_once(q, None).into_scalar_result().expect("q")),
        Some(s("Max"))
    );
    assert_eq!(
        pet_name(
            store
                .q_once_as_of(q, None, t1)
                .into_scalar_result()
                .expect("q")
        ),
        Some(s("Rex"))
    );
    let then = store
        .pull_as_of(alice, "[{:person/pet [:pet/name]}]", t1)
        .expect("pull as-of");
    assert_eq!(
        then.get(&kw!(:person/pet)),
        Some(&map(vec![
            (db_id.clone(), ref_(rex)),
            (kw!(:pet/name), s("Rex"))
        ]))
    );
}