Store module private variables in module, not in locals

Summary:
When loaded symbols are no longer reexported (D30534539), module private symbols (e. g. loaded symbols) are stored in local stack.

This diff changes privates to store them in module itself (like with public module variables).

The reasons for this diff:
* eventually replace `ValueRef` with explicit cell objects
* to implement loaded constants propagation

This diff it useful as is because it improves error message on trying to import private symbols. For example, for this code:

```
# a.bzl
x = 1

# b.bzl
load("a.bzl", "x")

# c.bzl
load("b.bzl", "x")
```

Previously error message was `Variable 'x' not found`, now error message is `Module symbol 'x' is not exported`.

Reviewed By: ndmitchell

Differential Revision: D30561707

fbshipit-source-id: 309948560569f48f00e520424edaf1a197c0ff70

Author: stepancheg

starlark/src/debug/evaluate.rs

@@ -49,7 +49,7 @@ impl<'v, 'a> Evaluator<'v, 'a> {
         // Push all the frozen variables into the module
         if let Some(frozen) = &self.module_variables {
             for (name, slot) in frozen.0.names.symbols() {
-                if let Some(value) = frozen.0.get_slot(*slot) {
+                if let Some(value) = frozen.0.get_slot(slot) {
                     self.module_env.set(name, value.to_value())
                 }
             }

starlark/src/environment/mod.rs

@@ -42,6 +42,10 @@ pub(crate) enum EnvironmentError {
     /// Cannot import private symbol, i.e. underscore prefixed
     #[error("Cannot import private symbol `{0}`")]
     CannotImportPrivateSymbol(String),
+    #[error("Module has no symbol `{0}`")]
+    ModuleHasNoSymbol(String),
+    #[error("Module symbol `{0}` is not exported")]
+    ModuleSymbolIsNotExported(String),
     #[error("No imports are available, you tried `{0}` (no call to `Evaluator.set_loader`)")]
     NoImportsAvailable(String),
 }

starlark/src/environment/modules.rs

@@ -26,6 +26,7 @@ use crate::{
         slots::{FrozenSlots, ModuleSlotId, MutableSlots},
         EnvironmentError,
     },
+    syntax::ast::Visibility,
     values::{
         docs,
         docs::{DocItem, DocString},
@@ -100,16 +101,25 @@ pub struct Module {
 }
 
 impl FrozenModule {
-    /// Get the value of the variable `name`.
-    /// Returns [`None`] if the variable isn't defined in the module or hasn't been set.
-    pub fn get(&self, name: &str) -> Option<OwnedFrozenValue> {
-        let slot = self.1.0.names.get_name(name)?;
+    /// Get value, exported or private by name.
+    pub(crate) fn get_any_visibility(&self, name: &str) -> Option<(OwnedFrozenValue, Visibility)> {
+        let (slot, vis) = self.1.0.names.get_name(name)?;
         // This code is safe because we know the frozen module ref keeps the values alive
         self.1
             .0
             .slots
             .get_slot(slot)
-            .map(|x| unsafe { OwnedFrozenValue::new(self.0.dupe(), x) })
+            .map(|x| (unsafe { OwnedFrozenValue::new(self.0.dupe(), x) }, vis))
+    }
+
+    /// Get the value of the exported variable `name`.
+    /// Returns [`None`] if the variable isn't defined in the module or it is private.
+    pub fn get(&self, name: &str) -> Option<OwnedFrozenValue> {
+        self.get_any_visibility(name)
+            .and_then(|(value, vis)| match vis {
+                Visibility::Private => None,
+                Visibility::Public => Some(value),
+            })
     }
 
     /// Iterate through all the names defined in this module.
@@ -157,7 +167,7 @@ impl FrozenModuleData {
     pub fn describe(&self) -> String {
         self.names
             .symbols()
-            .filter_map(|(name, slot)| Some((name, self.slots.get_slot(*slot)?)))
+            .filter_map(|(name, slot)| Some((name, self.slots.get_slot(slot)?)))
             .map(|(name, val)| val.to_value().describe(name))
             .join("\n")
     }
@@ -170,7 +180,7 @@ impl FrozenModuleData {
     /// Inefficient - only use in error paths.
     pub(crate) fn get_slot_name(&self, slot: ModuleSlotId) -> Option<String> {
         for (s, i) in self.names.symbols() {
-            if *i == slot {
+            if i == slot {
                 return Some(s.clone());
             }
         }
@@ -248,10 +258,21 @@ impl Module {
         unsafe { transmute!(&'v MutableSlots<'static>, &'v MutableSlots<'v>, &self.slots) }
     }
 
-    /// Get the value of the variable `name`, or [`None`] if the variable is not defined.
+    /// Get value, exported or private by name.
+    pub(crate) fn get_any_visibility<'v>(&'v self, name: &str) -> Option<(Value<'v>, Visibility)> {
+        let (slot, vis) = self.names.get_name(name)?;
+        let value = self.slots().get_slot(slot)?;
+        Some((value, vis))
+    }
+
+    /// Get the value of the exported variable `name`.
+    /// Returns [`None`] if the variable isn't defined in the module or it is private.
     pub fn get<'v>(&'v self, name: &str) -> Option<Value<'v>> {
-        let slot = self.names.get_name(name)?;
-        self.slots().get_slot(slot)
+        self.get_any_visibility(name)
+            .and_then(|(v, vis)| match vis {
+                Visibility::Private => None,
+                Visibility::Public => Some(v),
+            })
     }
 
     /// Freeze the environment, all its value will become immutable afterwards.
@@ -301,7 +322,7 @@ impl Module {
         self.frozen_heap.add_reference(&module.0);
         for (k, slot) in module.1.0.names.symbols() {
             if Self::is_public_symbol(k) {
-                if let Some(value) = module.1.0.slots.get_slot(*slot) {
+                if let Some(value) = module.1.0.slots.get_slot(slot) {
                     self.set(k, Value::new_frozen(value))
                 }
             }
@@ -316,9 +337,12 @@ impl Module {
         if !Self::is_public_symbol(symbol) {
             return Err(EnvironmentError::CannotImportPrivateSymbol(symbol.to_owned()).into());
         }
-        match module.get(symbol) {
-            None => Err(EnvironmentError::VariableNotFound(symbol.to_owned()).into()),
-            Some(v) => Ok(v.owned_value(self.frozen_heap())),
+        match module.get_any_visibility(symbol) {
+            None => Err(EnvironmentError::ModuleHasNoSymbol(symbol.to_owned()).into()),
+            Some((v, Visibility::Public)) => Ok(v.owned_value(self.frozen_heap())),
+            Some((_, Visibility::Private)) => {
+                Err(EnvironmentError::ModuleSymbolIsNotExported(symbol.to_owned()).into())
+            }
         }
     }
 

starlark/src/environment/names.rs

@@ -17,7 +17,7 @@
 
 use std::{cell::RefCell, collections::HashMap, iter::Iterator};
 
-use crate::environment::slots::ModuleSlotId;
+use crate::{environment::slots::ModuleSlotId, syntax::ast::Visibility};
 
 /// MutableNames are how we allocate slots (index-based) to variables
 /// (name-based). The slots field is the current active mapping of names to
@@ -37,10 +37,10 @@ use crate::environment::slots::ModuleSlotId;
 /// On an unscope, we do the reverse, putting things back to how they were
 /// before (apart from the total) number of slots required.
 #[derive(Debug)]
-pub(crate) struct MutableNames(RefCell<HashMap<String, ModuleSlotId>>);
+pub(crate) struct MutableNames(RefCell<HashMap<String, (ModuleSlotId, Visibility)>>);
 
 #[derive(Debug)]
-pub(crate) struct FrozenNames(HashMap<String, ModuleSlotId>);
+pub(crate) struct FrozenNames(HashMap<String, (ModuleSlotId, Visibility)>);
 
 impl MutableNames {
     pub fn new() -> Self {
@@ -54,37 +54,52 @@ impl MutableNames {
     /// Try and go back from a slot to a name.
     /// Inefficient - only use in error paths.
     pub fn get_slot(&self, slot: ModuleSlotId) -> Option<String> {
-        for (s, i) in &*self.0.borrow() {
+        for (s, (i, _vis)) in &*self.0.borrow() {
             if *i == slot {
                 return Some(s.clone());
             }
         }
         None
     }
 
-    pub fn get_name(&self, name: &str) -> Option<ModuleSlotId> {
+    pub(crate) fn get_name(&self, name: &str) -> Option<(ModuleSlotId, Visibility)> {
         self.0.borrow().get(name).copied()
     }
 
-    // Add a name, or if it's already there, return the existing name
-    pub fn add_name(&self, name: &str) -> ModuleSlotId {
+    /// Add a name with explicit visibility to the module.
+    pub(crate) fn add_name_visibility(&self, name: &str, vis: Visibility) -> ModuleSlotId {
         let mut x = self.0.borrow_mut();
-        match x.get(name) {
-            Some(v) => *v,
+        match x.get_mut(name) {
+            Some((slot, stored_vis)) => {
+                // Public visibility wins.
+                if *stored_vis == Visibility::Private {
+                    *stored_vis = vis;
+                }
+                *slot
+            }
             None => {
                 let slot = ModuleSlotId::new(x.len());
-                x.insert(name.to_owned(), slot);
+                x.insert(name.to_owned(), (slot, vis));
                 slot
             }
         }
     }
 
+    // Add an exported name, or if it's already there, return the existing name
+    pub fn add_name(&self, name: &str) -> ModuleSlotId {
+        self.add_name_visibility(name, Visibility::Public)
+    }
+
     pub fn hide_name(&self, name: &str) {
         self.0.borrow_mut().remove(name);
     }
 
     pub fn all_names(&self) -> HashMap<String, ModuleSlotId> {
-        self.0.borrow().clone()
+        self.0
+            .borrow()
+            .iter()
+            .map(|(name, (slot, _vis))| (name.to_owned(), *slot))
+            .collect()
     }
 
     pub fn freeze(self) -> FrozenNames {
@@ -93,11 +108,15 @@ impl MutableNames {
 }
 
 impl FrozenNames {
-    pub fn get_name(&self, name: &str) -> Option<ModuleSlotId> {
+    pub(crate) fn get_name(&self, name: &str) -> Option<(ModuleSlotId, Visibility)> {
         self.0.get(name).copied()
     }
 
-    pub fn symbols(&self) -> impl Iterator<Item = (&String, &ModuleSlotId)> {
-        self.0.iter()
+    /// Exported symbols.
+    pub fn symbols(&self) -> impl Iterator<Item = (&String, ModuleSlotId)> {
+        self.0.iter().flat_map(|(name, (slot, vis))| match vis {
+            Visibility::Private => None,
+            Visibility::Public => Some((name, *slot)),
+        })
     }
 }

starlark/src/eval/compiler/scope.rs

@@ -18,7 +18,7 @@
 use crate::{
     environment::{names::MutableNames, slots::ModuleSlotId},
     eval::runtime::slots::LocalSlotId,
-    syntax::ast::{AstAssign, AstStmt, Expr, Stmt, Visibility},
+    syntax::ast::{AstAssign, AstStmt, Expr, Stmt},
 };
 use std::collections::HashMap;
 
@@ -112,20 +112,12 @@ impl<'a> Scope<'a> {
     pub fn enter_module(module: &'a MutableNames, code: &AstStmt) -> Self {
         let mut locals = HashMap::new();
         Stmt::collect_defines(code, &mut locals);
-        let mut module_private = ScopeNames::default();
         for (x, vis) in locals {
-            match vis {
-                Visibility::Public => {
-                    module.add_name(x);
-                }
-                Visibility::Private => {
-                    module_private.add_name(x);
-                }
-            }
+            module.add_name_visibility(x, vis);
         }
         Self {
             module,
-            locals: vec![module_private],
+            locals: vec![ScopeNames::default()],
             unscopes: Vec::new(),
         }
     }
@@ -195,7 +187,9 @@ impl<'a> Scope<'a> {
                 return Some(Slot::Local(v));
             }
         }
-        self.module.get_name(name).map(Slot::Module)
+        self.module
+            .get_name(name)
+            .map(|(slot, _vis)| Slot::Module(slot))
     }
 
     pub fn get_name_or_panic(&mut self, name: &str) -> Slot {

starlark/src/eval/fragment/stmt.rs

@@ -502,7 +502,10 @@ impl Compiler<'_> {
                 let name = name.node;
                 let symbols = v.into_map(|(x, y)| {
                     (
-                        self.scope.get_name_or_panic(&x.node),
+                        match self.scope.get_name_or_panic(&x.node) {
+                            Slot::Local(..) => unreachable!("symbol need to be resolved to module"),
+                            Slot::Module(slot) => slot,
+                        },
                         y.node,
                         x.span.merge(y.span),
                     )
@@ -522,10 +525,7 @@ impl Compiler<'_> {
                             *span,
                             eval,
                         )?;
-                        match new_name {
-                            Slot::Local(slot) => eval.set_slot_local(*slot, value),
-                            Slot::Module(slot) => eval.set_slot_module(*slot, value),
-                        }
+                        eval.set_slot_module(*new_name, value)
                     }
                 })
             }

starlark/src/eval/tests.rs

@@ -1264,7 +1264,10 @@ fn test_load_reexport() {
     a.dialect_set(|d| d.enable_load_reexport = false);
     a.module("a", "x = 1");
     a.module("b", "load('a', 'x')");
-    a.fail("load('b', 'x')\nassert_eq(x, 1)", "Variable `x` not found");
+    a.fail(
+        "load('b', 'x')\nassert_eq(x, 1)",
+        "Module symbol `x` is not exported",
+    );
 }
 
 #[test]

starlark/src/syntax/ast.rs

@@ -180,7 +180,7 @@ pub enum AssignOp {
     RightShift,  // >>=
 }
 
-#[derive(Debug, Copy, Clone, Dupe)]
+#[derive(Debug, Copy, Clone, Dupe, Eq, PartialEq)]
 pub enum Visibility {
     Private,
     Public,