starlark/eval/compiler/
def.rs

1/*
2 * Copyright 2019 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
18//! Implementation of `def`.
19
20use std::cell::UnsafeCell;
21use std::collections::HashMap;
22use std::fmt;
23use std::fmt::Write;
24use std::ptr;
25
26use allocative::Allocative;
27use derivative::Derivative;
28use derive_more::Display;
29use dupe::Dupe;
30use once_cell::sync::Lazy;
31use starlark_derive::starlark_value;
32use starlark_derive::NoSerialize;
33use starlark_derive::VisitSpanMut;
34use starlark_map::StarlarkHasher;
35use starlark_syntax::eval_exception::EvalException;
36use starlark_syntax::slice_vec_ext::SliceExt;
37use starlark_syntax::syntax::def::DefParam;
38use starlark_syntax::syntax::def::DefParamIndices;
39use starlark_syntax::syntax::def::DefParamKind;
40use starlark_syntax::syntax::def::DefParams;
41
42use crate as starlark;
43use crate::any::ProvidesStaticType;
44use crate::codemap::CodeMap;
45use crate::codemap::Spanned;
46use crate::collections::Hashed;
47use crate::const_frozen_string;
48use crate::docs::DocFunction;
49use crate::docs::DocItem;
50use crate::docs::DocMember;
51use crate::docs::DocString;
52use crate::docs::DocStringKind;
53use crate::environment::FrozenModuleData;
54use crate::environment::Globals;
55use crate::eval::bc::bytecode::Bc;
56use crate::eval::bc::frame::alloca_frame;
57use crate::eval::compiler::def_inline::inline_def_body;
58use crate::eval::compiler::def_inline::InlineDefBody;
59use crate::eval::compiler::error::CompilerInternalError;
60use crate::eval::compiler::expr::ExprCompiled;
61use crate::eval::compiler::opt_ctx::OptCtx;
62use crate::eval::compiler::scope::payload::CstAssignIdent;
63use crate::eval::compiler::scope::payload::CstParameter;
64use crate::eval::compiler::scope::payload::CstPayload;
65use crate::eval::compiler::scope::payload::CstStmt;
66use crate::eval::compiler::scope::payload::CstTypeExpr;
67use crate::eval::compiler::scope::Captured;
68use crate::eval::compiler::scope::ScopeId;
69use crate::eval::compiler::span::IrSpanned;
70use crate::eval::compiler::stmt::OptimizeOnFreezeContext;
71use crate::eval::compiler::stmt::StmtCompileContext;
72use crate::eval::compiler::stmt::StmtsCompiled;
73use crate::eval::compiler::Compiler;
74use crate::eval::runtime::arguments::ArgumentsImpl;
75use crate::eval::runtime::arguments::ResolvedArgName;
76use crate::eval::runtime::evaluator::Evaluator;
77use crate::eval::runtime::frame_span::FrameSpan;
78use crate::eval::runtime::frozen_file_span::FrozenFileSpan;
79use crate::eval::runtime::params::spec::ParametersSpec;
80use crate::eval::runtime::profile::instant::ProfilerInstant;
81use crate::eval::runtime::slots::LocalSlotId;
82use crate::eval::runtime::slots::LocalSlotIdCapturedOrNot;
83use crate::eval::Arguments;
84use crate::starlark_complex_values;
85use crate::typing::callable_param::ParamIsRequired;
86use crate::typing::ParamSpec;
87use crate::typing::Ty;
88use crate::util::arc_str::ArcStr;
89use crate::values::frozen_ref::AtomicFrozenRefOption;
90use crate::values::function::FUNCTION_TYPE;
91use crate::values::typing::type_compiled::compiled::TypeCompiled;
92use crate::values::Freeze;
93use crate::values::FreezeResult;
94use crate::values::Freezer;
95use crate::values::FrozenHeap;
96use crate::values::FrozenRef;
97use crate::values::FrozenStringValue;
98use crate::values::FrozenValue;
99use crate::values::Heap;
100use crate::values::StarlarkValue;
101use crate::values::Trace;
102use crate::values::Tracer;
103use crate::values::Value;
104use crate::values::ValueLike;
105
106#[derive(thiserror::Error, Debug)]
107enum DefError {
108    #[error("Function has no type, while function was compiled with return type (internal error)")]
109    CheckReturnTypeNoType,
110}
111
112/// Store frozen `StmtCompiled`.
113/// This is initialized in `post_freeze`.
114struct StmtCompiledCell {
115    cell: UnsafeCell<Bc>,
116}
117
118unsafe impl<'v> Trace<'v> for StmtCompiledCell {
119    fn trace(&mut self, _: &Tracer<'v>) {
120        // Bytecode contains only frozen values.
121    }
122}
123
124unsafe impl Sync for StmtCompiledCell {}
125unsafe impl Send for StmtCompiledCell {}
126
127impl StmtCompiledCell {
128    fn new() -> StmtCompiledCell {
129        StmtCompiledCell {
130            cell: UnsafeCell::new(Bc::default()),
131        }
132    }
133
134    /// This function is unsafe if other thread is executing the stmt.
135    unsafe fn set(&self, value: Bc) {
136        ptr::drop_in_place(self.cell.get());
137        ptr::write(self.cell.get(), value);
138    }
139
140    fn get(&self) -> &Bc {
141        unsafe { &*self.cell.get() }
142    }
143}
144
145#[derive(Clone, Debug, VisitSpanMut)]
146pub(crate) struct ParameterName {
147    pub(crate) name: String,
148    captured: Captured,
149}
150
151#[derive(Clone, Debug, VisitSpanMut)]
152pub(crate) enum ParameterCompiled<T> {
153    Normal(
154        /// Name.
155        ParameterName,
156        /// Type.
157        Option<TypeCompiled<FrozenValue>>,
158        /// Default value.
159        Option<T>,
160    ),
161    Args(ParameterName, Option<TypeCompiled<FrozenValue>>),
162    KwArgs(ParameterName, Option<TypeCompiled<FrozenValue>>),
163}
164
165impl<T> ParameterCompiled<T> {
166    pub(crate) fn map_expr<U>(&self, f: impl FnMut(&T) -> U) -> ParameterCompiled<U> {
167        match self {
168            ParameterCompiled::Normal(n, o, t) => {
169                ParameterCompiled::Normal(n.clone(), *o, t.as_ref().map(f))
170            }
171            ParameterCompiled::Args(n, o) => ParameterCompiled::Args(n.clone(), *o),
172            ParameterCompiled::KwArgs(n, o) => ParameterCompiled::KwArgs(n.clone(), *o),
173        }
174    }
175
176    pub(crate) fn accepts_positional(&self) -> bool {
177        match self {
178            ParameterCompiled::Normal(..) => true,
179            _ => false,
180        }
181    }
182
183    pub(crate) fn captured(&self) -> Captured {
184        self.name_ty().0.captured
185    }
186
187    pub(crate) fn name_ty(&self) -> (&ParameterName, Option<TypeCompiled<FrozenValue>>) {
188        match self {
189            Self::Normal(n, t, _) => (n, *t),
190            Self::Args(n, t) => (n, *t),
191            Self::KwArgs(n, t) => (n, *t),
192        }
193    }
194
195    pub(crate) fn has_type(&self) -> bool {
196        match self.name_ty() {
197            (_, Some(_)) => true,
198            _ => false,
199        }
200    }
201
202    pub(crate) fn ty(&self) -> Ty {
203        match self.name_ty() {
204            (_, Some(t)) => t.as_ty().clone(),
205            (_, None) => Ty::any(),
206        }
207    }
208
209    pub(crate) fn required(&self) -> ParamIsRequired {
210        match self {
211            ParameterCompiled::Normal(_, _, None) => ParamIsRequired::Yes,
212            ParameterCompiled::Normal(_, _, Some(_)) => ParamIsRequired::No,
213            ParameterCompiled::Args(..) => ParamIsRequired::No,
214            ParameterCompiled::KwArgs(..) => ParamIsRequired::No,
215        }
216    }
217
218    pub(crate) fn is_star_or_star_star(&self) -> bool {
219        matches!(
220            self,
221            ParameterCompiled::Args(_, _) | ParameterCompiled::KwArgs(_, _)
222        )
223    }
224}
225
226#[derive(Debug, Clone, VisitSpanMut)]
227pub(crate) struct ParametersCompiled<T> {
228    pub(crate) params: Vec<IrSpanned<ParameterCompiled<T>>>,
229    pub(crate) indices: DefParamIndices,
230}
231
232impl<T> ParametersCompiled<T> {
233    /// How many expressions this parameters references (default values and types).
234    pub(crate) fn count_exprs(&self) -> u32 {
235        let mut count = 0;
236        for p in &self.params {
237            p.map_expr(|_e| count += 1);
238        }
239        count
240    }
241
242    /// How many parameter variables?
243    pub(crate) fn count_param_variables(&self) -> u32 {
244        self.params.len().try_into().unwrap()
245    }
246
247    /// Any parameter has type annotation?
248    pub(crate) fn has_types(&self) -> bool {
249        self.params.iter().any(|p| p.has_type())
250    }
251
252    /// Has `*args` or `*kwargs` parameter? `*` is fine.
253    pub(crate) fn has_args_or_kwargs(&self) -> bool {
254        self.params.iter().any(|p| {
255            matches!(
256                p.node,
257                ParameterCompiled::Args(..) | ParameterCompiled::KwArgs(..)
258            )
259        })
260    }
261
262    pub(crate) fn parameter_captures(&self) -> Vec<LocalSlotId> {
263        self.params
264            .iter()
265            .enumerate()
266            .filter_map(|(i, p)| {
267                if p.captured() == Captured::Yes {
268                    Some(LocalSlotId(i as u32))
269                } else {
270                    None
271                }
272            })
273            .collect()
274    }
275
276    pub(crate) fn to_ty_params(&self) -> ParamSpec {
277        ParamSpec::new_parts(
278            self.indices.pos_only().map(|i| {
279                let p = &self.params[i].node;
280                (p.required(), p.ty())
281            }),
282            self.indices.pos_or_named().map(|i| {
283                let p = &self.params[i].node;
284                (
285                    ArcStr::from(p.name_ty().0.name.as_str()),
286                    p.required(),
287                    p.ty(),
288                )
289            }),
290            self.indices.args.map(|i| {
291                let p = &self.params[i as usize].node;
292                p.ty()
293            }),
294            self.indices.named_only(self.params.len()).map(|i| {
295                let p = &self.params[i].node;
296                (
297                    ArcStr::from(p.name_ty().0.name.as_str()),
298                    p.required(),
299                    p.ty(),
300                )
301            }),
302            self.indices.kwargs.map(|i| {
303                let p = &self.params[i as usize].node;
304                p.ty()
305            }),
306        )
307        // TODO(nga): do not unwrap.
308        .unwrap()
309    }
310}
311
312/// Copy local variable slot to nested function.
313#[derive(Debug, Clone, Dupe)]
314pub(crate) struct CopySlotFromParent {
315    /// Slot in the outer function.
316    pub(crate) parent: LocalSlotIdCapturedOrNot,
317    /// Slot in the nested function.
318    pub(crate) child: LocalSlotIdCapturedOrNot,
319}
320
321/// Static info for `def`, `lambda` or module.
322#[derive(Derivative, Display)]
323#[derivative(Debug)]
324#[display("DefInfo")]
325pub(crate) struct DefInfo {
326    pub(crate) name: FrozenStringValue,
327    /// Span of function signature.
328    pub(crate) signature_span: FrozenFileSpan,
329    /// Indices of parameters, which are captured in nested defs.
330    parameter_captures: FrozenRef<'static, [LocalSlotId]>,
331    /// Type of this function, for the typechecker.
332    ty: Ty,
333    /// Codemap of the file where the function is declared.
334    pub(crate) codemap: FrozenRef<'static, CodeMap>,
335    /// The raw docstring pulled out of the AST.
336    pub(crate) docstring: Option<String>,
337    /// Slots this scope uses, including for parameters and `parent`.
338    /// Indexed by [`LocalSlotId`], values are variable names.
339    pub(crate) used: FrozenRef<'static, [FrozenStringValue]>,
340    /// Slots to copy from the parent.
341    /// Module-level identifiers are not copied over, to avoid excess copying.
342    pub(crate) parent: FrozenRef<'static, [CopySlotFromParent]>,
343    /// Statement compiled for non-frozen def.
344    #[derivative(Debug = "ignore")]
345    stmt_compiled: Bc,
346    // The compiled expression for the body of this definition, to be run
347    // after the parameters are evaluated.
348    #[derivative(Debug = "ignore")]
349    body_stmts: StmtsCompiled,
350    /// How to compile the statement on freeze.
351    stmt_compile_context: StmtCompileContext,
352    /// Function can be inlined.
353    pub(crate) inline_def_body: Option<InlineDefBody>,
354    /// Globals captured during function or module creation.
355    /// Only needed for debugger evaluation.
356    pub(crate) globals: FrozenRef<'static, Globals>,
357}
358
359impl DefInfo {
360    pub(crate) fn empty() -> FrozenRef<'static, DefInfo> {
361        static EMPTY: Lazy<DefInfo> = Lazy::new(|| DefInfo {
362            name: const_frozen_string!("<empty>"),
363            signature_span: FrozenFileSpan::default(),
364            parameter_captures: FrozenRef::new(&[]),
365            ty: Ty::any(),
366            codemap: FrozenRef::new(CodeMap::empty_static()),
367            docstring: None,
368            used: FrozenRef::new(&[]),
369            parent: FrozenRef::new(&[]),
370            stmt_compiled: Bc::default(),
371            body_stmts: StmtsCompiled::empty(),
372            stmt_compile_context: StmtCompileContext::default(),
373            inline_def_body: None,
374            globals: FrozenRef::new(Globals::empty()),
375        });
376        FrozenRef::new(&EMPTY)
377    }
378
379    pub(crate) fn for_module(
380        codemap: FrozenRef<'static, CodeMap>,
381        local_names: FrozenRef<'static, [FrozenStringValue]>,
382        parent: FrozenRef<'static, [CopySlotFromParent]>,
383        globals: FrozenRef<'static, Globals>,
384    ) -> DefInfo {
385        DefInfo {
386            name: const_frozen_string!("<module>"),
387            signature_span: FrozenFileSpan::default(),
388            parameter_captures: FrozenRef::new(&[]),
389            ty: Ty::any(),
390            codemap,
391            docstring: None,
392            used: local_names,
393            parent,
394            stmt_compiled: Bc::default(),
395            body_stmts: StmtsCompiled::empty(),
396            stmt_compile_context: StmtCompileContext::default(),
397            inline_def_body: None,
398            globals,
399        }
400    }
401}
402
403#[derive(Clone, Debug, VisitSpanMut)]
404pub(crate) struct DefCompiled {
405    pub(crate) function_name: String,
406    pub(crate) params: ParametersCompiled<IrSpanned<ExprCompiled>>,
407    pub(crate) return_type: Option<TypeCompiled<FrozenValue>>,
408    pub(crate) info: FrozenRef<'static, DefInfo>,
409}
410
411impl Compiler<'_, '_, '_, '_> {
412    fn parameter_name(&mut self, ident: &CstAssignIdent) -> ParameterName {
413        let binding_id = ident.payload.expect("no binding for parameter");
414        let binding = self.scope_data.get_binding(binding_id);
415        ParameterName {
416            name: ident.node.ident.clone(),
417            captured: binding.captured,
418        }
419    }
420
421    fn parameter(
422        &mut self,
423        x: &Spanned<DefParam<'_, CstPayload>>,
424    ) -> Result<IrSpanned<ParameterCompiled<IrSpanned<ExprCompiled>>>, CompilerInternalError> {
425        let span = FrameSpan::new(FrozenFileSpan::new(self.codemap, x.span));
426        let parameter_name = self.parameter_name(x.ident);
427        Ok(IrSpanned {
428            span,
429            node: match &x.node.kind {
430                DefParamKind::Regular(_mode, default_value) => ParameterCompiled::Normal(
431                    parameter_name,
432                    self.expr_for_type(x.ty).map(|t| t.node),
433                    default_value.as_ref().map(|d| self.expr(d)).transpose()?,
434                ),
435                DefParamKind::Args => ParameterCompiled::Args(
436                    parameter_name,
437                    self.expr_for_type(x.ty).map(|t| t.node),
438                ),
439                DefParamKind::Kwargs => ParameterCompiled::KwArgs(
440                    parameter_name,
441                    self.expr_for_type(x.ty).map(|t| t.node),
442                ),
443            },
444        })
445    }
446
447    pub fn function(
448        &mut self,
449        name: &str,
450        signature_span: FrozenFileSpan,
451        scope_id: ScopeId,
452        params: &[CstParameter],
453        return_type: Option<&CstTypeExpr>,
454        suite: &CstStmt,
455    ) -> Result<ExprCompiled, CompilerInternalError> {
456        let file = self.codemap.file_span(suite.span);
457        let function_name = format!("{}.{}", file.file.filename(), name);
458        let name = self.eval.frozen_heap().alloc_str_intern(name);
459
460        let DefParams { params, indices } = match DefParams::unpack(params, &self.codemap) {
461            Ok(def_params) => def_params,
462            Err(e) => return Err(CompilerInternalError::from_eval_exception(e)),
463        };
464
465        // The parameters run in the scope of the parent, so compile them with the outer
466        // scope
467        let params: Vec<_> = params
468            .iter()
469            .map(|x| self.parameter(x))
470            .collect::<Result<_, CompilerInternalError>>()?;
471        let params = ParametersCompiled { params, indices };
472        let return_type = self.expr_for_type(return_type).map(|t| t.node);
473
474        let ty = Ty::function(
475            params.to_ty_params(),
476            return_type.map_or(Ty::any(), |t| t.as_ty().clone()),
477        );
478
479        self.enter_scope(scope_id);
480
481        let docstring = DocString::extract_raw_starlark_docstring(suite);
482        let body = self.stmt(suite, false)?;
483        let scope_id = self.exit_scope();
484        let scope_names = self.scope_data.get_scope(scope_id);
485
486        let has_types = return_type.is_some() || params.has_types();
487
488        let inline_def_body = if has_types {
489            // It is harder to inline if a function declares parameter types or return type.
490            None
491        } else {
492            inline_def_body(&params, &body)
493        };
494
495        let param_count = params.count_param_variables();
496
497        let used = self.eval.frozen_heap().alloc_any_slice(&scope_names.used);
498        let info = self.eval.module_env.frozen_heap().alloc_any(DefInfo {
499            name,
500            signature_span,
501            parameter_captures: self
502                .eval
503                .frozen_heap()
504                .alloc_any_slice(&params.parameter_captures()),
505            ty,
506            codemap: self.codemap,
507            docstring,
508            used,
509            parent: self.eval.frozen_heap().alloc_any_slice(&scope_names.parent),
510            stmt_compiled: body.as_bc(
511                &self.compile_context(return_type.is_some()),
512                used,
513                param_count,
514                self.eval.module_env.frozen_heap(),
515            ),
516            body_stmts: body,
517            inline_def_body,
518            stmt_compile_context: self.compile_context(return_type.is_some()),
519            globals: self.globals,
520        });
521
522        Ok(ExprCompiled::Def(DefCompiled {
523            function_name,
524            params,
525            return_type,
526            info,
527        }))
528    }
529}
530
531/// Starlark function internal representation and implementation of
532/// [`StarlarkValue`].
533#[derive(Derivative, NoSerialize, ProvidesStaticType, Trace, Allocative)]
534#[derivative(Debug)]
535pub(crate) struct DefGen<V> {
536    pub(crate) parameters: ParametersSpec<V>, // The parameters, **kwargs etc including defaults (which are evaluated afresh each time)
537    /// Indices of parameters, which are captured in nested defs.
538    /// This is a copy of `DefInfo.parameter_captures`.
539    parameter_captures: FrozenRef<'static, [LocalSlotId]>,
540    // The types of the parameters.
541    // (Sparse indexed array, (0, argm T) implies parameter 0 named arg must have type T).
542    parameter_types: Vec<(LocalSlotId, String, TypeCompiled<FrozenValue>)>,
543    pub(crate) return_type: Option<TypeCompiled<FrozenValue>>, // The return type annotation for the function
544    /// Data created during function compilation but before function instantiation.
545    /// `DefInfo` can be shared by multiple `def` instances, for example,
546    /// `lambda` functions can be instantiated multiple times.
547    pub(crate) def_info: FrozenRef<'static, DefInfo>,
548    /// Any variables captured from the outer scope (nested def/lambda).
549    /// Values are either [`Value`] or [`FrozenValue`] pointing respectively to
550    /// [`ValueCaptured`] or [`FrozenValueCaptured`].
551    captured: Vec<V>,
552    // Important to ignore these field as it probably references DefGen in a cycle
553    #[derivative(Debug = "ignore")]
554    /// A reference to the module where the function is defined after the module has been frozen.
555    /// When the module is not frozen yet, this field contains `None`, and function's module
556    /// can be accessed from evaluator's module.
557    #[allocative(skip)]
558    pub(crate) module: AtomicFrozenRefOption<FrozenModuleData>,
559    /// This field is only used in `FrozenDef`. It is populated in `post_freeze`.
560    #[derivative(Debug = "ignore")]
561    #[allocative(skip)]
562    optimized_on_freeze_stmt: StmtCompiledCell,
563}
564
565impl<V> Display for DefGen<V> {
566    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
567        write!(f, "{}", self.parameters.signature())
568    }
569}
570
571pub(crate) type Def<'v> = DefGen<Value<'v>>;
572pub(crate) type FrozenDef = DefGen<FrozenValue>;
573
574starlark_complex_values!(Def);
575
576impl<'v> Def<'v> {
577    pub(crate) fn new(
578        parameters: ParametersSpec<Value<'v>>,
579        parameter_types: Vec<(LocalSlotId, String, TypeCompiled<FrozenValue>)>,
580        return_type: Option<TypeCompiled<FrozenValue>>,
581        stmt: FrozenRef<'static, DefInfo>,
582        eval: &mut Evaluator<'v, '_, '_>,
583    ) -> anyhow::Result<Value<'v>> {
584        let captured = stmt
585            .parent
586            .as_ref()
587            .map(|copy| eval.clone_slot_capture(copy, &stmt));
588        Ok(eval.heap().alloc(Self {
589            parameters,
590            parameter_captures: stmt.parameter_captures,
591            parameter_types,
592            return_type,
593            captured,
594            module: AtomicFrozenRefOption::new(eval.top_frame_def_frozen_module(false)?),
595            optimized_on_freeze_stmt: StmtCompiledCell::new(),
596            def_info: stmt,
597        }))
598    }
599}
600
601impl<'v> Freeze for Def<'v> {
602    type Frozen = FrozenDef;
603
604    fn freeze(self, freezer: &Freezer) -> FreezeResult<Self::Frozen> {
605        let parameters = self.parameters.freeze(freezer)?;
606        let parameter_types = self.parameter_types.freeze(freezer)?;
607        let return_type = self.return_type.freeze(freezer)?;
608        let captured = self.captured.try_map(|x| x.freeze(freezer))?;
609        let module = AtomicFrozenRefOption::new(self.module.load_relaxed());
610        Ok(FrozenDef {
611            parameters,
612            parameter_captures: self.parameter_captures,
613            parameter_types,
614            return_type,
615            def_info: self.def_info,
616            captured,
617            module,
618            optimized_on_freeze_stmt: self.optimized_on_freeze_stmt,
619        })
620    }
621}
622
623pub(crate) trait DefLike<'v> {
624    const FROZEN: bool;
625}
626
627impl<'v> DefLike<'v> for DefGen<Value<'v>> {
628    const FROZEN: bool = false;
629}
630
631impl<'v> DefLike<'v> for DefGen<FrozenValue> {
632    const FROZEN: bool = true;
633}
634
635#[starlark_value(type = FUNCTION_TYPE)]
636impl<'v, V: ValueLike<'v>> StarlarkValue<'v> for DefGen<V>
637where
638    Self: ProvidesStaticType<'v> + DefLike<'v>,
639{
640    fn name_for_call_stack(&self, _me: Value<'v>) -> String {
641        self.def_info.name.as_str().to_owned()
642    }
643
644    fn invoke(
645        &self,
646        me: Value<'v>,
647        args: &Arguments<'v, '_>,
648        eval: &mut Evaluator<'v, '_, '_>,
649    ) -> crate::Result<Value<'v>> {
650        self.invoke_impl(me, &args.0, eval)
651    }
652
653    fn documentation(&self) -> DocItem {
654        let mut parameter_types = vec![Ty::any(); self.parameters.len()];
655        for (idx, _, ty) in &self.parameter_types {
656            // Local slot number for parameter is the same as parameter index.
657            parameter_types[idx.0 as usize] = ty.as_ty().clone();
658        }
659
660        let return_type = self.return_type.map_or(Ty::any(), |r| r.as_ty().clone());
661
662        let function_docs = DocFunction::from_docstring(
663            DocStringKind::Starlark,
664            self.parameters
665                .documentation(parameter_types, HashMap::new()),
666            return_type,
667            self.def_info.docstring.as_ref().map(String::as_ref),
668        );
669
670        DocItem::Member(DocMember::Function(function_docs))
671    }
672
673    fn typechecker_ty(&self) -> Option<Ty> {
674        Some(self.def_info.ty.clone())
675    }
676
677    fn write_hash(&self, hasher: &mut StarlarkHasher) -> crate::Result<()> {
678        // It's hard to come up with a good hash here, but let's at least make an effort.
679        self.def_info.name.write_hash(hasher)
680    }
681}
682
683impl<'v, V: ValueLike<'v>> DefGen<V>
684where
685    Self: DefLike<'v>,
686{
687    pub(crate) fn bc(&self) -> &Bc {
688        if Self::FROZEN {
689            self.optimized_on_freeze_stmt.get()
690        } else {
691            &self.def_info.stmt_compiled
692        }
693    }
694
695    fn check_parameter_types(&self, eval: &mut Evaluator<'v, '_, '_>) -> crate::Result<()> {
696        let start = if eval.typecheck_profile.enabled {
697            Some(ProfilerInstant::now())
698        } else {
699            None
700        };
701        for (i, arg_name, ty) in &self.parameter_types {
702            match eval.current_frame.get_slot(i.to_captured_or_not()) {
703                None => {
704                    panic!("Not allowed optional unassigned with type annotations on them")
705                }
706                Some(v) => ty.check_type(v, Some(arg_name))?,
707            }
708        }
709        if let Some(start) = start {
710            eval.typecheck_profile
711                .add(self.def_info.name, start.elapsed());
712        }
713        Ok(())
714    }
715
716    pub(crate) fn check_return_type(
717        &self,
718        ret: Value<'v>,
719        eval: &mut Evaluator<'v, '_, '_>,
720    ) -> crate::Result<()> {
721        let return_type_ty: TypeCompiled<FrozenValue> = self
722            .return_type
723            .ok_or_else(|| crate::Error::new_other(DefError::CheckReturnTypeNoType))?;
724        let start = if eval.typecheck_profile.enabled {
725            Some(ProfilerInstant::now())
726        } else {
727            None
728        };
729        return_type_ty.check_type(ret, None)?;
730        if let Some(start) = start {
731            eval.typecheck_profile
732                .add(self.def_info.name, start.elapsed());
733        }
734        Ok(())
735    }
736
737    #[inline(always)]
738    fn invoke_impl<'a, A: ArgumentsImpl<'v, 'a>>(
739        &self,
740        me: Value<'v>,
741        args: &A,
742        eval: &mut Evaluator<'v, '_, '_>,
743    ) -> crate::Result<Value<'v>>
744    where
745        'v: 'a,
746    {
747        let bc = self.bc();
748        alloca_frame(
749            eval,
750            bc.local_count,
751            bc.max_stack_size,
752            bc.max_loop_depth,
753            |eval| {
754                // SAFETY: `slots` is unique: `alloca_frame` just allocated the frame,
755                //   so there are no references to the frame except `eval.current_frame`.
756                //   We use `slots` only in `collect_inline`,
757                //   which does not have access to `eval` thus cannot access the frame indirectly.
758                let slots = unsafe { eval.current_frame.locals_mut() };
759                self.parameters.collect_inline(args, slots, eval.heap())?;
760                self.invoke_raw(me, eval)
761            },
762        )
763    }
764
765    pub(crate) fn invoke_with_args<'a, A: ArgumentsImpl<'v, 'a>>(
766        &self,
767        me: Value<'v>,
768        args: &A,
769        eval: &mut Evaluator<'v, '_, '_>,
770    ) -> crate::Result<Value<'v>>
771    where
772        'v: 'a,
773    {
774        // This is trivial function which delegates to `invoke_impl`.
775        // `invoke_impl` is called from two places,
776        // giving this function different name makes this function easier to see in profiler.
777        self.invoke_impl(me, args, eval)
778    }
779
780    /// Invoke the function, assuming that:
781    /// * the frame has been allocated and stored in `eval.current_frame`
782    /// * the arguments have been collected into the frame
783    #[inline(always)]
784    fn invoke_raw(
785        &self,
786        me: Value<'v>,
787        eval: &mut Evaluator<'v, '_, '_>,
788    ) -> crate::Result<Value<'v>> {
789        // println!("invoking {}", self.def.stmt.name.node);
790
791        if !self.parameter_types.is_empty() {
792            self.check_parameter_types(eval)?;
793        }
794
795        // Parameters are collected into local slots without captures
796        // (to avoid even more branches in parameter capture),
797        // and this loop wraps captured parameters.
798        for &captured in &*self.parameter_captures {
799            eval.wrap_local_slot_captured(captured);
800        }
801
802        // Copy over the parent slots.
803        // Explicitly check `self.captured` is not empty to avoid accessing
804        // self.def_info.scope_names which is two indirections.
805        if !self.captured.is_empty() {
806            for (copy, captured) in self.def_info.parent.iter().zip(self.captured.iter()) {
807                eval.current_frame.set_slot(copy.child, captured.to_value());
808            }
809        }
810
811        if Self::FROZEN {
812            debug_assert!(self.module.load_relaxed().is_some());
813        }
814
815        eval.eval_bc(me, self.bc())
816            .map_err(EvalException::into_error)
817    }
818
819    pub(crate) fn resolve_arg_name(&self, name: Hashed<&str>) -> ResolvedArgName {
820        self.parameters.resolve_name(name)
821    }
822
823    pub(crate) fn dump_debug(&self) -> String {
824        let mut w = String::new();
825        writeln!(w, "Bytecode:").unwrap();
826        self.bc()
827            .dump_debug()
828            .lines()
829            .for_each(|l| writeln!(w, "  {}", l).unwrap());
830        w
831    }
832}
833
834impl FrozenDef {
835    pub(crate) fn post_freeze(
836        &self,
837        module: FrozenRef<FrozenModuleData>,
838        heap: &Heap,
839        frozen_heap: &FrozenHeap,
840    ) {
841        // Module passed to this function is not always module where the function is declared:
842        // A function can be created in a frozen module and frozen later in another module.
843        // `def_module` variable contains a module where this `def` is declared.
844        let def_module = match self.module.load_relaxed() {
845            None => {
846                self.module.store_relaxed(module);
847                module
848            }
849            Some(module) => module,
850        };
851
852        // Now perform the optimization of function body with fully frozen module:
853        // all module variables are frozen, so we can inline more aggressively.
854        let body_optimized = self
855            .def_info
856            .body_stmts
857            .optimize(&mut OptCtx::new(
858                &mut OptimizeOnFreezeContext {
859                    module: def_module.as_ref(),
860                    heap,
861                    frozen_heap,
862                },
863                self.parameters.len().try_into().unwrap(),
864            ))
865            .as_bc(
866                &self.def_info.stmt_compile_context,
867                self.def_info.used,
868                self.parameters.len() as u32,
869                frozen_heap,
870            );
871
872        // Store the optimized body.
873        // This is (relatively) safe because we know that during freeze
874        // nobody has a reference to stmt: nobody is executing this `def`.
875        unsafe {
876            self.optimized_on_freeze_stmt.set(body_optimized);
877        }
878    }
879}
starlark/eval/compiler/def.rs

starlark/eval/compiler/
def.rs
