{Python,Ruby,JS}: Add class_id information to ExternalInstance valu…

…es; use for in-VM IsA checks (#1472)

Create `Constants` type which wraps two HashMaps describing Symbol ->
Term and ID -> Symbol relationships. This second ID -> Symbol index is
populated by reading the optional `class_id` information in
`ExternalInstance` and will be used to implement in-core IsA checks
where present.

Extend `vm#isa` to attempt an in-core check when the `ExternalInstance` value
we're processing has the necessary `class_id` information declared. At time of
commit Python is the only language which is capabale of passing this information
to the core.

* {JS,Python,Ruby} Format ExternalInstance values with user-provided class names if present  (#1488)

Co-authored-by: Gabe Jackson <17556281+gj@users.noreply.github.com>

languages/js/src/Host.ts

@@ -456,12 +456,29 @@ export class Host<Query, Resource> {
       }
       default: {
         let instanceId: number | undefined = undefined;
-        if (isConstructor(v)) instanceId = this.getType(v)?.id;
+        let classId: number | undefined = undefined;
 
         // pass a string class repr *for registered types only*, otherwise pass
         // undefined (allow core to differentiate registered or not)
         const v_cast = v as NullishOrHasConstructor;
-        let classRepr: string | undefined = v_cast?.constructor?.name;
+        let classRepr: string | undefined = undefined;
+
+        if (isConstructor(v)) {
+          instanceId = this.getType(v)?.id;
+          classId = instanceId;
+          classRepr = this.getType(v)?.name;
+        } else {
+          const v_constructor: Class | undefined = v_cast?.constructor;
+
+          // pass classId for instances of *registered classes* only
+          if (v_constructor !== undefined && this.types.has(v_constructor)) {
+            classId = this.getType(v_constructor)?.id;
+            classRepr = this.getType(v_constructor)?.name;
+          }
+        }
+
+        // pass classRepr for *registered* classes only, pass undefined
+        // otherwise
         if (classRepr !== undefined && !this.types.has(classRepr)) {
           classRepr = undefined;
         }
@@ -474,6 +491,7 @@ export class Host<Query, Resource> {
               constructor: undefined,
               repr: repr(v),
               class_repr: classRepr,
+              class_id: classId,
             },
           },
         };

languages/js/src/Polar.ts

@@ -171,7 +171,6 @@ export class Polar<Query, Resource> {
 
   // Register MROs, load Polar code, and check inline queries.
   private async loadSources(sources: Source[]): Promise<void> {
-    this.getHost().registerMros();
     this.#ffiPolar.load(sources);
     this.processMessages();
     return this.checkInlineQueries();
@@ -244,6 +243,7 @@ export class Polar<Query, Resource> {
   registerClass(cls: Class, params?: ClassParams): void {
     const clsName = this.getHost().cacheClass(cls, params);
     this.registerConstant(cls, clsName);
+    this.getHost().registerMros();
   }
 
   /**

languages/js/src/types.ts

@@ -180,6 +180,7 @@ interface PolarInstance {
     instance_id: number;
     repr: string;
     class_repr?: string;
+    class_id?: number;
     constructor?: PolarTerm;
   };
 }
@@ -691,6 +692,7 @@ export interface ClassParams {
    * `name` property.
    */
   name?: string;
+
   /**
    * A Map or object with string keys containing types for fields. Used for
    * data filtering.

languages/python/oso/polar/host.py

@@ -294,23 +294,30 @@ def to_polar(self, v):
         #       ]
         #   }
         else:
-            instance_id = None
             import inspect
 
+            instance_id = None
+            class_id = None
+
             # maintain consistent IDs for registered classes
             if inspect.isclass(v):
                 if v in self.types:
-                    instance_id = self.types[v].id
+                    class_id = instance_id = self.types[v].id
 
             # pass the class_repr only for registered types otherwise None
-            class_repr = type(v).__name__
-            class_repr = class_repr if class_repr in self.types else None
+            class_repr = self.types[type(v)].name if type(v) in self.types else None
+
+            # pass class_id for classes & instances of registered classes,
+            # otherwise pass None
+            if type(v) in self.types:
+                class_id = self.types[type(v)].id
 
             val = {
                 "ExternalInstance": {
                     "instance_id": self.cache_instance(v, instance_id),
                     "repr": None,
                     "class_repr": class_repr,
+                    "class_id": class_id,
                 }
             }
         term = {"value": val}

languages/python/oso/polar/polar.py

@@ -95,7 +95,6 @@ def load_str(self, string: str):
 
     # Register MROs, load Polar code, and check inline queries.
     def _load_sources(self, sources: List[Source]):
-        self.host.register_mros()
         self.ffi_polar.load(sources)
         self.check_inline_queries()
 
@@ -241,12 +240,14 @@ def register_class(
             data filtering.
         """
         # TODO: let's add example usage here or at least a proper docstring for the arguments
-        cls_name = self.host.cache_class(
+
+        name = self.host.cache_class(
             cls,
             name=name,
             fields=fields,
         )
-        self.register_constant(cls, cls_name)
+        self.register_constant(cls, name)
+        self.host.register_mros()
 
     def register_constant(self, value, name):
         """

languages/ruby/lib/oso/polar/host.rb

@@ -319,17 +319,24 @@ def to_polar(value) # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplex
                   end
                 else
                   instance_id = nil
-                  instance_id = types[value].id if value.is_a?(Class) && types.key?(value)
-
-                  # only pass class_repr for registered types
+                  class_id = nil
                   class_repr = nil
-                  class_repr = value.class.to_s if types.key?(value.class)
+                  # id=class_id,
+
+                  # pass `class_id` & `class_repr` for registered types
+                  if value.is_a?(Class) && types.key?(value)
+                    instance_id = class_id = types[value].id
+                  elsif types.key?(value.class)
+                    class_id = types[value.class].id
+                    class_repr = types[value.class].name
+                  end
 
                   {
                     'ExternalInstance' => {
                       'instance_id' => cache_instance(value, id: instance_id),
                       'repr' => nil,
-                      'class_repr' => class_repr
+                      'class_repr' => class_repr,
+                      'class_id' => class_id
                     }
                   }
                 end

languages/ruby/lib/oso/polar/polar.rb

@@ -217,6 +217,7 @@ def register_class(cls, name: nil, fields: nil, combine_query: nil, build_query:
           exec_query: exec_query || maybe_mtd(cls, :exec_query)
         )
         register_constant(cls, name: name)
+        host.register_mros
       end
 
       # Register a Ruby object with Polar.
@@ -336,7 +337,6 @@ def maybe_mtd(cls, mtd)
       # Register MROs, load Polar code, and check inline queries.
       # @param sources [Array<Source>] Polar sources to load.
       def load_sources(sources)
-        host.register_mros
         ffi_polar.load(sources)
         check_inline_queries
       end

languages/rust/oso/src/host/value.rs

@@ -113,6 +113,7 @@ This may mean you performed an operation in your policy over an unbound variable
                     instance_id: id,
                     repr: Some(std::any::type_name::<Self>().to_owned()),
                     class_repr: Some(std::any::type_name::<Self>().to_owned()),
+                    class_id: None,
                 })
             }
             PolarValue::List(l) => {

polar-core/src/constants.rs

@@ -0,0 +1,40 @@
+use crate::terms::{Symbol, Term};
+use std::collections::HashMap;
+
+#[derive(Default, Debug)]
+pub(crate) struct Constants {
+    // Symbol -> Term (populated by *all* constants)
+    pub symbol_to_term: HashMap<Symbol, Term>,
+    // Symbol -> class_id (populated by class constants)
+    class_symbol_to_id: HashMap<Symbol, u64>,
+    // class_id -> Symbol (populated by class constants)
+    class_id_to_symbol: HashMap<u64, Symbol>,
+}
+
+impl Constants {
+    pub(crate) fn insert(&mut self, name: Symbol, value: Term) {
+        self.symbol_to_term.insert(name, value);
+    }
+
+    pub(crate) fn insert_class(&mut self, name: Symbol, value: Term, class_id: u64) {
+        self.insert(name.clone(), value);
+        self.class_symbol_to_id.insert(name.clone(), class_id);
+        self.class_id_to_symbol.insert(class_id, name);
+    }
+
+    pub(crate) fn contains_key(&self, name: &Symbol) -> bool {
+        self.symbol_to_term.contains_key(name)
+    }
+
+    pub(crate) fn get(&self, name: &Symbol) -> Option<&Term> {
+        self.symbol_to_term.get(name)
+    }
+
+    pub(crate) fn get_class_id_for_symbol(&self, symbol: &Symbol) -> Option<&u64> {
+        self.class_symbol_to_id.get(symbol)
+    }
+
+    pub(crate) fn get_symbol_for_class_id(&self, id: &u64) -> Option<&Symbol> {
+        self.class_id_to_symbol.get(id)
+    }
+}

polar-core/src/folder.rs

@@ -147,6 +147,7 @@ pub fn fold_external_instance<T: Folder>(
         constructor,
         repr,
         class_repr,
+        class_id,
     }: ExternalInstance,
     fld: &mut T,
 ) -> ExternalInstance {
@@ -155,6 +156,7 @@ pub fn fold_external_instance<T: Folder>(
         constructor: constructor.map(|t| fld.fold_term(t)),
         repr: repr.map(|r| fld.fold_string(r)),
         class_repr: class_repr.map(|r| fld.fold_string(r)),
+        class_id: class_id.map(|id| fld.fold_instance_id(id)),
     }
 }
 
@@ -271,6 +273,7 @@ mod tests {
             constructor: None,
             repr: None,
             class_repr: None,
+            class_id: None,
         }));
         let instance_pattern = term!(value!(Pattern::Instance(InstanceLiteral {
             tag: sym!("d"),

polar-core/src/kb.rs

@@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 pub use super::bindings::Bindings;
+use super::constants::Constants;
 use super::counter::Counter;
 use super::diagnostic::Diagnostic;
 use super::error::{invalid_state, PolarError, PolarResult, RuntimeError, ValidationError};
@@ -26,9 +27,9 @@ impl RuleParamMatch {
 pub struct KnowledgeBase {
     /// A map of bindings: variable name → value. The VM uses a stack internally,
     /// but can translate to and from this type.
-    constants: Bindings,
+    constants: Constants,
     /// Map of class name -> MRO list where the MRO list is a list of class instance IDs
-    mro: HashMap<Symbol, Vec<u64>>,
+    pub mro: HashMap<Symbol, Vec<u64>>,
 
     /// Map from contents to filename for files loaded into the KB.
     loaded_content: HashMap<String, String>,
@@ -574,7 +575,26 @@ impl KnowledgeBase {
             }
             .into());
         }
-        self.constants.insert(name, value);
+
+        if let Value::ExternalInstance(ExternalInstance {
+            class_id,
+            instance_id,
+            ..
+        }) = *value.value()
+        {
+            if class_id.map_or(false, |id| id == instance_id) {
+                // ExternalInstance values with matching class_id & instance_id represent *classes*
+                // whose class_id we want to index for later type checking & MRO resolution
+                self.constants.insert_class(name, value, instance_id)
+            } else {
+                // ExternalInstance values with differing `class_id` and
+                // `instance_id` represent *instances* of classes whose class_id
+                // should not be registered
+                self.constants.insert(name, value)
+            }
+        } else {
+            self.constants.insert(name, value)
+        }
         Ok(())
     }
 
@@ -585,7 +605,15 @@ impl KnowledgeBase {
 
     /// Getter for `constants` map without exposing it for mutation.
     pub fn get_registered_constants(&self) -> &Bindings {
-        &self.constants
+        &self.constants.symbol_to_term
+    }
+
+    pub(crate) fn get_symbol_for_class_id(&self, id: &u64) -> Option<&Symbol> {
+        self.constants.get_symbol_for_class_id(id)
+    }
+
+    pub(crate) fn get_class_id_for_symbol(&self, symbol: &Symbol) -> Option<&u64> {
+        self.constants.get_class_id_for_symbol(symbol)
     }
 
     // TODO(gj): currently no way to distinguish classes from other registered constants in the
@@ -879,6 +907,7 @@ mod tests {
                     constructor: None,
                     repr: None,
                     class_repr: None,
+                    class_id: None,
                 })),
             )
             .unwrap();
@@ -1321,6 +1350,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1331,6 +1361,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1341,6 +1372,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1401,6 +1433,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1411,6 +1444,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1421,6 +1455,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();
@@ -1432,6 +1467,7 @@ mod tests {
                 constructor: None,
                 repr: None,
                 class_repr: None,
+                class_id: None,
             })),
         )
         .unwrap();

polar-core/src/lib.rs

@@ -8,6 +8,7 @@ extern crate maplit;
 pub mod macros;
 
 mod bindings;
+mod constants;
 mod counter;
 pub mod data_filtering;
 mod debugger;

polar-core/src/macros.rs

@@ -241,6 +241,7 @@ impl From<u64> for ExternalInstance {
             constructor: None,
             repr: None,
             class_repr: None,
+            class_id: None,
         }
     }
 }