starlark/environment/
globals.rs

1/*
2 * Copyright 2018 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18use std::sync::Arc;
19
20use allocative::Allocative;
21use dupe::Dupe;
22use itertools::Itertools;
23use once_cell::sync::Lazy;
24use once_cell::sync::OnceCell;
25
26use crate::__derive_refs::components::NativeCallableComponents;
27use crate::collections::symbol::map::SymbolMap;
28use crate::collections::SmallMap;
29use crate::docs::DocItem;
30use crate::docs::DocModule;
31use crate::docs::DocString;
32use crate::docs::DocStringKind;
33use crate::docs::DocType;
34use crate::stdlib;
35pub use crate::stdlib::LibraryExtension;
36use crate::typing::Ty;
37use crate::values::function::NativeFunc;
38use crate::values::function::SpecialBuiltinFunction;
39use crate::values::namespace::value::MaybeDocHiddenValue;
40use crate::values::namespace::FrozenNamespace;
41use crate::values::types::function::NativeFunction;
42use crate::values::AllocFrozenValue;
43use crate::values::FrozenHeap;
44use crate::values::FrozenHeapRef;
45use crate::values::FrozenStringValue;
46use crate::values::FrozenValue;
47use crate::values::Value;
48
49/// The global values available during execution.
50#[derive(Clone, Dupe, Debug, Allocative)]
51pub struct Globals(Arc<GlobalsData>);
52
53type GlobalValue = MaybeDocHiddenValue<'static, FrozenValue>;
54
55#[derive(Debug, Allocative)]
56struct GlobalsData {
57    heap: FrozenHeapRef,
58    variables: SymbolMap<GlobalValue>,
59    variable_names: Vec<FrozenStringValue>,
60    docstring: Option<String>,
61}
62
63/// Used to build a [`Globals`] value.
64#[derive(Debug)]
65pub struct GlobalsBuilder {
66    // The heap everything is allocated in
67    heap: FrozenHeap,
68    // Normal top-level variables, e.g. True/hash
69    variables: SymbolMap<GlobalValue>,
70    // The list of struct fields, pushed to the end
71    namespace_fields: Vec<SmallMap<FrozenStringValue, GlobalValue>>,
72    /// The raw docstring for this module
73    ///
74    /// FIXME(JakobDegen): This should probably be removed. Having a docstring on a `GlobalsBuilder`
75    /// doesn't really make sense, because there's no way good way to combine multiple docstrings.
76    docstring: Option<String>,
77}
78
79impl Globals {
80    /// Create an empty [`Globals`], with no functions in scope.
81    pub fn new() -> Self {
82        GlobalsBuilder::new().build()
83    }
84
85    /// Create a [`Globals`] following the
86    /// [Starlark standard](https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions).
87    pub fn standard() -> Self {
88        GlobalsBuilder::standard().build()
89    }
90
91    /// Create a [`Globals`] combining those functions in the Starlark standard plus
92    /// all those defined in [`LibraryExtension`].
93    ///
94    /// This function is public to use in the `starlark` binary,
95    /// but users of starlark should list the extensions they want explicitly.
96    #[doc(hidden)]
97    pub fn extended_internal() -> Self {
98        GlobalsBuilder::extended().build()
99    }
100
101    /// Empty globals.
102    pub(crate) fn empty() -> &'static Globals {
103        static EMPTY: Lazy<Globals> = Lazy::new(|| GlobalsBuilder::new().build());
104        &EMPTY
105    }
106
107    /// Create a [`Globals`] combining those functions in the Starlark standard plus
108    /// all those given in the [`LibraryExtension`] arguments.
109    pub fn extended_by(extensions: &[LibraryExtension]) -> Self {
110        GlobalsBuilder::extended_by(extensions).build()
111    }
112
113    /// This function is only safe if you first call `heap` and keep a reference to it.
114    /// Therefore, don't expose it on the public API.
115    pub(crate) fn get<'v>(&'v self, name: &str) -> Option<Value<'v>> {
116        self.get_frozen(name).map(FrozenValue::to_value)
117    }
118
119    /// This function is only safe if you first call `heap` and keep a reference to it.
120    /// Therefore, don't expose it on the public API.
121    pub(crate) fn get_frozen(&self, name: &str) -> Option<FrozenValue> {
122        self.0.variables.get_str(name).map(|x| x.value)
123    }
124
125    /// Get all the names defined in this environment.
126    pub fn names(&self) -> impl Iterator<Item = FrozenStringValue> + '_ {
127        self.0.variable_names.iter().copied()
128    }
129
130    /// Iterate over all the items in this environment.
131    /// Note returned values are owned by this globals.
132    pub fn iter(&self) -> impl Iterator<Item = (&str, FrozenValue)> {
133        self.0.variables.iter().map(|(n, v)| (n.as_str(), v.value))
134    }
135
136    pub(crate) fn heap(&self) -> &FrozenHeapRef {
137        &self.0.heap
138    }
139
140    /// Print information about the values in this object.
141    pub fn describe(&self) -> String {
142        self.0
143            .variables
144            .iter()
145            .map(|(name, val)| val.value.to_value().describe(name.as_str()))
146            .join("\n")
147    }
148
149    /// Get the documentation for the object itself
150    pub fn docstring(&self) -> Option<&str> {
151        self.0.docstring.as_deref()
152    }
153
154    /// Get the documentation for both the object itself, and its members.
155    pub fn documentation(&self) -> DocModule {
156        let (docs, members) = common_documentation(
157            &self.0.docstring,
158            self.0
159                .variables
160                .iter()
161                .filter(|(_, v)| !v.doc_hidden)
162                .map(|(n, v)| (n.as_str(), v.value)),
163        );
164        DocModule {
165            docs,
166            members: members.collect(),
167        }
168    }
169}
170
171impl GlobalsBuilder {
172    /// Create an empty [`GlobalsBuilder`], with no functions in scope.
173    pub fn new() -> Self {
174        Self {
175            heap: FrozenHeap::new(),
176            variables: SymbolMap::new(),
177            namespace_fields: Vec::new(),
178            docstring: None,
179        }
180    }
181
182    /// Create a [`GlobalsBuilder`] following the
183    /// [Starlark standard](https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions).
184    pub fn standard() -> Self {
185        stdlib::standard_environment()
186    }
187
188    /// Create a [`GlobalsBuilder`] combining those functions in the Starlark standard plus
189    /// all those defined in [`LibraryExtension`].
190    pub(crate) fn extended() -> Self {
191        Self::extended_by(LibraryExtension::all())
192    }
193
194    /// Create a [`GlobalsBuilder`] combining those functions in the Starlark standard plus
195    /// all those defined in [`LibraryExtension`].
196    pub fn extended_by(extensions: &[LibraryExtension]) -> Self {
197        let mut res = Self::standard();
198        for x in extensions {
199            x.add(&mut res);
200        }
201        res
202    }
203
204    /// Add a nested namespace to the builder. If `f` adds the definition `foo`,
205    /// it will end up on a namespace `name`, accessible as `name.foo`.
206    pub fn namespace(&mut self, name: &str, f: impl FnOnce(&mut GlobalsBuilder)) {
207        self.namespace_inner(name, false, f)
208    }
209
210    /// Same as `namespace`, but this value will not show up in generated documentation.
211    pub fn namespace_no_docs(&mut self, name: &str, f: impl FnOnce(&mut GlobalsBuilder)) {
212        self.namespace_inner(name, true, f)
213    }
214
215    fn namespace_inner(
216        &mut self,
217        name: &str,
218        doc_hidden: bool,
219        f: impl FnOnce(&mut GlobalsBuilder),
220    ) {
221        self.namespace_fields.push(SmallMap::new());
222        f(self);
223        let fields = self.namespace_fields.pop().unwrap();
224        self.set_inner(
225            name,
226            self.heap.alloc(FrozenNamespace::new(fields)),
227            doc_hidden,
228        );
229    }
230
231    /// A fluent API for modifying [`GlobalsBuilder`] and returning the result.
232    pub fn with(mut self, f: impl FnOnce(&mut Self)) -> Self {
233        f(&mut self);
234        self
235    }
236
237    /// A fluent API for modifying [`GlobalsBuilder`] using [`namespace`](GlobalsBuilder::namespace).
238    pub fn with_namespace(mut self, name: &str, f: impl Fn(&mut GlobalsBuilder)) -> Self {
239        self.namespace(name, f);
240        self
241    }
242
243    /// Called at the end to build a [`Globals`].
244    pub fn build(self) -> Globals {
245        let mut variable_names: Vec<_> = self
246            .variables
247            .keys()
248            .map(|x| self.heap.alloc_str_intern(x.as_str()))
249            .collect();
250        variable_names.sort();
251        Globals(Arc::new(GlobalsData {
252            heap: self.heap.into_ref(),
253            variables: self.variables,
254            variable_names,
255            docstring: self.docstring,
256        }))
257    }
258
259    /// Set a value in the [`GlobalsBuilder`].
260    pub fn set<'v, V: AllocFrozenValue>(&'v mut self, name: &str, value: V) {
261        let value = value.alloc_frozen_value(&self.heap);
262        self.set_inner(name, value, false)
263    }
264
265    fn set_inner<'v>(&'v mut self, name: &str, value: FrozenValue, doc_hidden: bool) {
266        let value = MaybeDocHiddenValue {
267            value,
268            doc_hidden,
269            phantom: Default::default(),
270        };
271        match self.namespace_fields.last_mut() {
272            None => {
273                // TODO(nga): do not quietly ignore redefinitions.
274                self.variables.insert(name, value)
275            }
276            Some(fields) => {
277                let name = self.heap.alloc_str(name);
278                fields.insert(name, value)
279            }
280        };
281    }
282
283    /// Set a method. This function is usually called from code
284    /// generated by `starlark_derive` and rarely needs to be called manually.
285    pub fn set_function<F>(
286        &mut self,
287        name: &str,
288        components: NativeCallableComponents,
289        as_type: Option<(Ty, DocType)>,
290        ty: Option<Ty>,
291        special_builtin_function: Option<SpecialBuiltinFunction>,
292        f: F,
293    ) where
294        F: NativeFunc,
295    {
296        self.set(
297            name,
298            NativeFunction {
299                function: Box::new(f),
300                name: name.to_owned(),
301                speculative_exec_safe: components.speculative_exec_safe,
302                as_type: as_type.as_ref().map(|x| x.0.dupe()),
303                ty: ty.unwrap_or_else(|| {
304                    Ty::from_native_callable_components(
305                        &components,
306                        as_type.as_ref().map(|x| x.0.dupe()),
307                    )
308                    .unwrap() // TODO(nga): do not unwrap.
309                }),
310                docs: components.into_docs(as_type),
311                special_builtin_function,
312            },
313        )
314    }
315
316    /// Heap where globals are allocated. Can be used to allocate additional values.
317    pub fn frozen_heap(&self) -> &FrozenHeap {
318        &self.heap
319    }
320
321    /// Allocate a value using the same underlying heap as the [`GlobalsBuilder`],
322    /// only intended for values that are referred to by those which are passed
323    /// to [`set`](GlobalsBuilder::set).
324    pub fn alloc<'v, V: AllocFrozenValue>(&'v self, value: V) -> FrozenValue {
325        value.alloc_frozen_value(&self.heap)
326    }
327
328    /// Set per module docstring.
329    ///
330    /// This function is called by the `starlark_derive` generated code
331    /// and rarely needs to be called manually.
332    pub fn set_docstring(&mut self, docstring: &str) {
333        self.docstring = Some(docstring.to_owned());
334    }
335}
336
337/// Used to create globals.
338pub struct GlobalsStatic(OnceCell<Globals>);
339
340impl GlobalsStatic {
341    /// Create a new [`GlobalsStatic`].
342    pub const fn new() -> Self {
343        Self(OnceCell::new())
344    }
345
346    fn globals(&'static self, x: impl FnOnce(&mut GlobalsBuilder)) -> &'static Globals {
347        self.0.get_or_init(|| GlobalsBuilder::new().with(x).build())
348    }
349
350    /// Get a function out of the object. Requires that the function passed only set a single
351    /// value. If populated via a `#[starlark_module]`, that means a single function in it.
352    pub fn function(&'static self, x: impl FnOnce(&mut GlobalsBuilder)) -> FrozenValue {
353        let globals = self.globals(x);
354        assert!(
355            globals.0.variables.len() == 1,
356            "GlobalsBuilder.function must have exactly 1 member, you had {}",
357            globals
358                .names()
359                .map(|s| format!("`{}`", s.as_str()))
360                .join(", ")
361        );
362
363        globals.0.variables.values().next().unwrap().value
364    }
365
366    /// Move all the globals in this [`GlobalsBuilder`] into a new one. All variables will
367    /// only be allocated once (ensuring things like function comparison works properly).
368    pub fn populate(&'static self, x: impl FnOnce(&mut GlobalsBuilder), out: &mut GlobalsBuilder) {
369        let globals = self.globals(x);
370        for (name, value) in globals.0.variables.iter() {
371            out.set_inner(name.as_str(), value.value, value.doc_hidden)
372        }
373        out.docstring = globals.0.docstring.clone();
374    }
375}
376
377pub(crate) fn common_documentation<'a>(
378    docstring: &Option<String>,
379    members: impl IntoIterator<Item = (&'a str, FrozenValue)>,
380) -> (Option<DocString>, impl Iterator<Item = (String, DocItem)>) {
381    let main_docs = docstring
382        .as_ref()
383        .and_then(|ds| DocString::from_docstring(DocStringKind::Rust, ds));
384    let member_docs = members
385        .into_iter()
386        .map(|(name, val)| (name.to_owned(), val.to_value().documentation()))
387        .sorted_by(|(l, _), (r, _)| Ord::cmp(l, r));
388
389    (main_docs, member_docs)
390}
391
392#[cfg(test)]
393mod tests {
394    use starlark_derive::starlark_module;
395
396    use super::*;
397    use crate as starlark;
398
399    #[test]
400    fn test_send_sync()
401    where
402        Globals: Send + Sync,
403    {
404    }
405
406    #[starlark_module]
407    fn register_foo(builder: &mut GlobalsBuilder) {
408        fn foo() -> anyhow::Result<i32> {
409            Ok(1)
410        }
411    }
412
413    #[test]
414    fn test_doc_hidden() {
415        let mut globals = GlobalsBuilder::new();
416        globals.namespace_no_docs("ns_hidden", |_| {});
417        globals.namespace("ns", |globals| {
418            globals.namespace_no_docs("nested_ns_hidden", |_| {});
419            globals.set("x", FrozenValue::new_none());
420        });
421        let docs = globals.build().documentation();
422
423        let (k, v) = docs.members.into_iter().exactly_one().ok().unwrap();
424        assert_eq!(&k, "ns");
425        let DocItem::Module(docs) = v else {
426            unreachable!()
427        };
428        assert_eq!(&docs.members.into_keys().exactly_one().ok().unwrap(), "x");
429    }
430}
starlark/environment/globals.rs

starlark/environment/
globals.rs
