WIP: nested pull.

query-parser/src/parse.rs

@@ -67,6 +67,7 @@ use self::mentat_query::{
     Order,
     OrJoin,
     OrWhereClause,
+    NamedPullAttribute,
     NotJoin,
     Pattern,
     PatternNonValuePlace,
@@ -191,6 +192,7 @@ def_parser!(Query, order, Order, {
 def_matches_plain_symbol!(Query, the, "the");
 def_matches_plain_symbol!(Query, pull, "pull");
 def_matches_plain_symbol!(Query, wildcard, "*");
+def_matches_keyword!(Query, alias_as, "as");
 
 pub struct Where<'a>(std::marker::PhantomData<&'a ()>);
 
@@ -303,11 +305,28 @@ def_parser!(Query, aggregate, Aggregate, {
          })
 });
 
+def_parser!(Query, pull_concrete_attribute_ident, PullConcreteAttribute, {
+    forward_keyword().map(|k| PullConcreteAttribute::Ident(::std::rc::Rc::new(k.clone())))
+});
+
+def_parser!(Query, pull_aliased_attribute, PullAttributeSpec, {
+    vector().of_exactly(
+        (Query::pull_concrete_attribute_ident()
+            .skip(Query::alias_as()),
+        forward_keyword().map(|alias| Some(::std::rc::Rc::new(alias.clone()))))
+            .map(|(attribute, alias)|
+                PullAttributeSpec::Attribute(
+                    NamedPullAttribute { attribute, alias })))
+});
+
 def_parser!(Query, pull_concrete_attribute, PullAttributeSpec, {
-    forward_keyword().map(|k|
-        PullAttributeSpec::Attribute(
-            PullConcreteAttribute::Ident(
-                ::std::rc::Rc::new(k.clone()))))
+    Query::pull_concrete_attribute_ident()
+        .map(|attribute|
+            PullAttributeSpec::Attribute(
+                NamedPullAttribute {
+                    attribute,
+                    alias: None,
+                }))
 });
 
 def_parser!(Query, pull_wildcard_attribute, PullAttributeSpec, {
@@ -316,6 +335,7 @@ def_parser!(Query, pull_wildcard_attribute, PullAttributeSpec, {
 
 def_parser!(Query, pull_attribute, PullAttributeSpec, {
     choice([
+        try(Query::pull_aliased_attribute()),
         try(Query::pull_concrete_attribute()),
         try(Query::pull_wildcard_attribute()),
         // TODO: reversed keywords, entids (with aliases, presumably…).
@@ -346,7 +366,11 @@ fn validate_attributes<'a, I>(attrs: I) -> std::result::Result<(), &'static str>
                     return Err("wildcard with specified attributes");
                 }
             },
-            // TODO: map form.
+            &PullAttributeSpec::Nested(ref _attr, ref patterns) => {
+                if patterns.is_empty() {
+                    return Err("empty nested pull map");
+                }
+            },
         }
     }
     Ok(())
@@ -1205,23 +1229,27 @@ mod test {
 
         let foo_bar = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "bar"));
         let foo_baz = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "baz"));
+        let foo_horse = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "horse"));
         assert_edn_parses_to!(Query::pull_concrete_attribute,
                               ":foo/bar",
                               PullAttributeSpec::Attribute(
-                                  PullConcreteAttribute::Ident(foo_bar.clone())));
+                                  PullConcreteAttribute::Ident(foo_bar.clone()).into()));
         assert_edn_parses_to!(Query::pull_attribute,
                               ":foo/bar",
                               PullAttributeSpec::Attribute(
-                                  PullConcreteAttribute::Ident(foo_bar.clone())));
+                                  PullConcreteAttribute::Ident(foo_bar.clone()).into()));
         assert_edn_parses_to!(Find::elem,
-                              "(pull ?v [:foo/bar :foo/baz])",
+                              "(pull ?v [[:foo/bar :as :foo/horse] :foo/baz])",
                               Element::Pull(Pull {
                                   var: Variable::from_valid_name("?v"),
                                   patterns: vec![
                                       PullAttributeSpec::Attribute(
-                                          PullConcreteAttribute::Ident(foo_bar.clone())),
+                                          NamedPullAttribute {
+                                              attribute: PullConcreteAttribute::Ident(foo_bar.clone()),
+                                              alias: Some(foo_horse),
+                                          }),
                                       PullAttributeSpec::Attribute(
-                                          PullConcreteAttribute::Ident(foo_baz.clone())),
+                                          PullConcreteAttribute::Ident(foo_baz.clone()).into()),
                                   ],
                               }));
         assert_parse_failure_contains!(Find::elem,
@@ -1242,7 +1270,7 @@ mod test {
                             PullAttributeSpec::Attribute(
                                 PullConcreteAttribute::Ident(
                                     ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "bar"))
-                                )
+                                ).into()
                             ),
                         ] })]),
                 where_clauses: vec![

query-pull/src/errors.rs

@@ -22,6 +22,16 @@ error_chain! {
             description("unnamed attribute")
             display("attribute {:?} has no name", id)
         }
+
+        RepeatedDbId {
+            description(":db/id repeated")
+            display(":db/id repeated")
+        }
+
+        NonRefNestedPullAttribute {
+            description("nested pull attribute is non-ref")
+            display("nested pull attribute is non-ref")
+        }
     }
 
     links {

query-pull/src/lib.rs

@@ -79,18 +79,22 @@ use std::iter::{
 };
 
 use mentat_core::{
+    Binding,
     Cloned,
     Entid,
     HasSchema,
     NamespacedKeyword,
     Schema,
     StructuredMap,
+    TypedValue,
     ValueRc,
+    ValueType,
 };
 
 use mentat_db::cache;
 
 use mentat_query::{
+    NamedPullAttribute,
     PullAttributeSpec,
     PullConcreteAttribute,
 };
@@ -110,7 +114,7 @@ pub fn pull_attributes_for_entity<A>(schema: &Schema,
                                      attributes: A) -> Result<StructuredMap>
     where A: IntoIterator<Item=Entid> {
     let attrs = attributes.into_iter()
-                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
+                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e).into()))
                           .collect();
     Puller::prepare(schema, attrs)?
         .pull(schema, db, once(entity))
@@ -130,7 +134,7 @@ pub fn pull_attributes_for_entities<E, A>(schema: &Schema,
     where E: IntoIterator<Item=Entid>,
           A: IntoIterator<Item=Entid> {
     let attrs = attributes.into_iter()
-                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
+                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e).into()))
                           .collect();
     Puller::prepare(schema, attrs)?
         .pull(schema, db, entities)
@@ -141,17 +145,21 @@ pub struct Puller {
     // The domain of this map is the set of attributes to fetch.
     // The range is the set of aliases to use in the output.
     attributes: BTreeMap<Entid, ValueRc<NamespacedKeyword>>,
+
+    // The original spec for this puller.
     attribute_spec: cache::AttributeSpec,
+
+    // If :db/id is mentioned in the attribute list, its alias is this.
+    db_id_alias: Option<ValueRc<NamespacedKeyword>>,
+
+    // A pull expression can be arbitrarily nested. We represent this both
+    // within the `attribute_spec` itself and also as a nested set of `Puller`s.
+    // When an attribute in the list above returns an entity -- and it should! --
+    // it is accumulated and we recurse down into these nested layers.
+    nested: BTreeMap<Entid, Puller>,
 }
 
 impl Puller {
-    pub fn prepare_simple_attributes(schema: &Schema, attributes: Vec<Entid>) -> Result<Puller> {
-        Puller::prepare(schema,
-                        attributes.into_iter()
-                                  .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
-                                  .collect())
-    }
-
     pub fn prepare(schema: &Schema, attributes: Vec<PullAttributeSpec>) -> Result<Puller> {
         // TODO: eventually this entry point will handle aliasing and that kind of
         // thing. For now it's just a convenience.
@@ -159,12 +167,15 @@ impl Puller {
         let lookup_name = |i: &Entid| {
             // In the unlikely event that we have an attribute with no name, we bail.
             schema.get_ident(*i)
-                    .map(|ident| ValueRc::new(ident.clone()))
-                    .ok_or_else(|| ErrorKind::UnnamedAttribute(*i))
+                  .map(|ident| ValueRc::new(ident.clone()))
+                  .ok_or_else(|| ErrorKind::UnnamedAttribute(*i))
         };
 
         let mut names: BTreeMap<Entid, ValueRc<NamespacedKeyword>> = Default::default();
         let mut attrs: BTreeSet<Entid> = Default::default();
+        let db_id = ::std::rc::Rc::new(NamespacedKeyword::new("db", "id"));
+        let mut db_id_alias = None;
+
         for attr in attributes.iter() {
             match attr {
                 &PullAttributeSpec::Wildcard => {
@@ -175,15 +186,48 @@ impl Puller {
                     }
                     break;
                 },
-                &PullAttributeSpec::Attribute(PullConcreteAttribute::Ident(ref i)) => {
-                    if let Some(entid) = schema.get_entid(i) {
-                        names.insert(entid.into(), i.to_value_rc());
-                        attrs.insert(entid.into());
+                &PullAttributeSpec::Attribute(NamedPullAttribute {
+                    ref attribute,
+                    ref alias,
+                }) => {
+                    let alias = alias.as_ref()
+                                     .map(|ref r| r.to_value_rc());
+                    match attribute {
+                        // Handle :db/id.
+                        &PullConcreteAttribute::Ident(ref i) if i.as_ref() == db_id.as_ref() => {
+                            // We only allow :db/id once.
+                            if db_id_alias.is_some() {
+                                bail!(ErrorKind::RepeatedDbId);
+                            }
+                            db_id_alias = Some(alias.unwrap_or_else(|| db_id.to_value_rc()));
+                        },
+                        &PullConcreteAttribute::Ident(ref i) => {
+                            if let Some(entid) = schema.get_entid(i) {
+                                let name = alias.unwrap_or_else(|| i.to_value_rc());
+                                names.insert(entid.into(), name);
+                                attrs.insert(entid.into());
+                            }
+                        },
+                        &PullConcreteAttribute::Entid(ref entid) => {
+                            let name = alias.map(Ok).unwrap_or_else(|| lookup_name(entid))?;
+                            names.insert(*entid, name);
+                            attrs.insert(*entid);
+                        },
                     }
                 },
-                &PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(ref entid)) => {
-                    names.insert(*entid, lookup_name(entid)?);
-                    attrs.insert(*entid);
+
+                // An attribute that nests must be ref-typed.
+                &PullAttributeSpec::Nested(ref attribute, ref patterns) => {
+                    let value_type = attribute.get_attribute(schema)
+                                              .map(|(a, _e)| a.value_type);
+                    let is_ref_typed = value_type.map(|v| v == ValueType::Ref)
+                                                 .unwrap_or(false);
+                    if !is_ref_typed {
+                        bail!(ErrorKind::NonRefNestedPullAttribute);
+                    }
+
+                    // TODO
+                    unimplemented!();
                 },
             }
         }
@@ -191,6 +235,8 @@ impl Puller {
         Ok(Puller {
             attributes: names,
             attribute_spec: cache::AttributeSpec::specified(&attrs, schema),
+            db_id_alias,
+            nested: Default::default(),
         })
     }
 
@@ -205,9 +251,7 @@ impl Puller {
         // - Recursing. (TODO: we'll need AttributeCaches to not overwrite in case of recursion! And
         //   ideally not do excess work when some entity/attribute pairs are known.)
         // - Building a structure by walking the pull expression with the caches.
-        //   TODO: aliases.
         // TODO: limits.
-        // TODO: fts.
 
         // Build a cache for these attributes and entities.
         // TODO: use the store's existing cache!
@@ -222,6 +266,17 @@ impl Puller {
         // TODO: should we walk `e` then `a`, or `a` then `e`? Possibly the right answer
         // is just to collect differently!
         let mut maps = BTreeMap::new();
+
+        // Collect :db/id if requested.
+        if let Some(ref alias) = self.db_id_alias {
+            for e in entities.iter() {
+                let mut r = maps.entry(*e)
+                                .or_insert(ValueRc::new(StructuredMap::default()));
+                let mut m = ValueRc::get_mut(r).unwrap();
+                m.insert(alias.clone(), Binding::Scalar(TypedValue::Ref(*e)));
+            }
+        }
+
         for (name, cache) in self.attributes.iter().filter_map(|(a, name)|
             caches.forward_attribute_cache_for_attribute(schema, *a)
                   .map(|cache| (name.clone(), cache))) {

query/src/lib.rs

@@ -55,7 +55,11 @@ pub use edn::{
 };
 
 use mentat_core::{
+    Attribute,
     FromRc,
+    HasSchema,
+    KnownEntid,
+    Schema,
     TypedValue,
     ValueRc,
     ValueType,
@@ -499,12 +503,39 @@ pub enum PullConcreteAttribute {
     Entid(i64),
 }
 
+impl PullConcreteAttribute {
+    pub fn get_attribute<'s>(&self, schema: &'s Schema) -> Option<(&'s Attribute, KnownEntid)> {
+        match self {
+            &PullConcreteAttribute::Ident(ref rc) => {
+                schema.attribute_for_ident(rc.as_ref())
+            },
+            &PullConcreteAttribute::Entid(e) => {
+                schema.attribute_for_entid(e).map(|a| (a, KnownEntid(e)))
+            },
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct NamedPullAttribute {
+    pub attribute: PullConcreteAttribute,
+    pub alias: Option<Rc<NamespacedKeyword>>,
+}
+
+impl From<PullConcreteAttribute> for NamedPullAttribute {
+    fn from(a: PullConcreteAttribute) -> Self {
+        NamedPullAttribute {
+            attribute: a,
+            alias: None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum PullAttributeSpec {
     Wildcard,
-    Attribute(PullConcreteAttribute),
-    // PullMapSpec(Vec<…>),
-    // AttributeWithOpts(PullConcreteAttribute, …),
+    Attribute(NamedPullAttribute),
+    Nested(PullConcreteAttribute, Vec<PullAttributeSpec>),
     // LimitedAttribute(PullConcreteAttribute, u64),  // Limit nil => Attribute instead.
     // DefaultedAttribute(PullConcreteAttribute, PullDefaultValue),
 }
@@ -522,14 +553,32 @@ impl std::fmt::Display for PullConcreteAttribute {
     }
 }
 
+impl std::fmt::Display for NamedPullAttribute {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        if let &Some(ref alias) = &self.alias {
+            write!(f, "{} :as {}", self.attribute, alias)
+        } else {
+            write!(f, "{}", self.attribute)
+        }
+    }
+}
+
+
 impl std::fmt::Display for PullAttributeSpec {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             &PullAttributeSpec::Wildcard => {
                 write!(f, "*")
             },
-            &PullAttributeSpec::Attribute(ref a) => {
-                write!(f, "{}", a)
+            &PullAttributeSpec::Attribute(ref attr) => {
+                write!(f, "{}", attr)
+            },
+            &PullAttributeSpec::Nested(ref attr, ref patterns) => {
+                write!(f, "{{{} [", attr)?;
+                for p in patterns {
+                    write!(f, " {}", p)?;
+                }
+                write!(f, "]}}")
             },
         }
     }

tests/pull.rs

@@ -116,7 +116,9 @@ fn test_simple_pull() {
     assert_eq!(pulled, expected);
 
     // Now test pull inside the query itself.
-    let query = r#"[:find ?hood (pull ?district [:district/name :district/region])
+    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"])
@@ -127,22 +129,24 @@ fn test_simple_pull() {
                                            .into_rel_result()
                                            .expect("results");
 
-    let beacon_district: Vec<(NamespacedKeyword, TypedValue)> = vec![
-        (kw!(:district/name), "Greater Duwamish".into()),
+    let beacon_district_pull: Vec<(NamespacedKeyword, TypedValue)> = vec![
+        (kw!(:db/id), TypedValue::Ref(beacon_district)),
+        (kw!(:district/district), "Greater Duwamish".into()),
         (kw!(:district/region), schema.get_entid(&NamespacedKeyword::new("region", "se")).unwrap().into())
     ];
-    let beacon_district: StructuredMap = beacon_district.into();
-    let capitol_district: Vec<(NamespacedKeyword, TypedValue)> = vec![
-        (kw!(:district/name), "East".into()),
+    let beacon_district_pull: StructuredMap = beacon_district_pull.into();
+    let capitol_district_pull: Vec<(NamespacedKeyword, TypedValue)> = vec![
+        (kw!(:db/id), TypedValue::Ref(capitol_district)),
+        (kw!(:district/district), "East".into()),
         (kw!(:district/region), schema.get_entid(&NamespacedKeyword::new("region", "e")).unwrap().into())
     ];
-    let capitol_district: StructuredMap = capitol_district.into();
+    let capitol_district_pull: StructuredMap = capitol_district_pull.into();
 
     let expected = RelResult {
                        width: 2,
                        values: vec![
-                           TypedValue::Ref(capitol).into(), capitol_district.into(),
-                           TypedValue::Ref(beacon).into(), beacon_district.into(),
+                           TypedValue::Ref(capitol).into(), capitol_district_pull.into(),
+                           TypedValue::Ref(beacon).into(), beacon_district_pull.into(),
                        ].into(),
                    };
     assert_eq!(results, expected.clone());
@@ -158,14 +162,19 @@ fn test_simple_pull() {
 
     // 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 [:neighborhood/name]) . :in ?hood
+    let query = r#"[:find (pull ?hood [[:db/id :as :neighborhood/id]
+                                       :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"))].into();
+    let expected: StructuredMap = vec![
+        (kw!(:neighborhood/name), TypedValue::from("Beacon Hill")),
+        (kw!(:neighborhood/id), TypedValue::Ref(beacon)),
+    ].into();
     assert_eq!(result, expected.into());
 
     // Collect the names and regions of all districts.