starlark/eval/bc/compiler/
expr.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//! Compile expressions.
19
20use std::collections::HashSet;
21
22use starlark_syntax::slice_vec_ext::SliceExt;
23
24use crate::collections::Hashed;
25use crate::collections::SmallMap;
26use crate::eval::bc::compiler::if_compiler::write_if_else;
27use crate::eval::bc::instr_impl::*;
28use crate::eval::bc::slow_arg::BcInstrSlowArg;
29use crate::eval::bc::stack_ptr::BcSlot;
30use crate::eval::bc::stack_ptr::BcSlotIn;
31use crate::eval::bc::stack_ptr::BcSlotInRange;
32use crate::eval::bc::stack_ptr::BcSlotOut;
33use crate::eval::bc::writer::BcWriter;
34use crate::eval::compiler::expr::Builtin1;
35use crate::eval::compiler::expr::Builtin2;
36use crate::eval::compiler::expr::CompareOp;
37use crate::eval::compiler::expr::ExprCompiled;
38use crate::eval::compiler::expr::ExprLogicalBinOp;
39use crate::eval::compiler::expr::MaybeNot;
40use crate::eval::compiler::span::IrSpanned;
41use crate::eval::runtime::frame_span::FrameSpan;
42use crate::values::layout::value_not_special::FrozenValueNotSpecial;
43use crate::values::FrozenStringValue;
44use crate::values::FrozenValue;
45use crate::values::ValueLike;
46
47/// Try extract consecutive definitely initialized locals from expressions.
48fn try_slot_range<'a>(
49    exprs: impl IntoIterator<Item = &'a IrSpanned<ExprCompiled>>,
50    bc: &BcWriter,
51) -> Option<BcSlotInRange> {
52    let mut range = BcSlotInRange::default();
53    for expr in exprs {
54        let local = expr.as_local_non_captured()?;
55        let slot = bc.try_definitely_assigned(local)?;
56        if !range.try_push(slot) {
57            return None;
58        }
59    }
60    Some(range)
61}
62
63/// Compile several expressions into consecutive registers.
64pub(crate) fn write_exprs<'a>(
65    exprs: impl IntoIterator<Item = &'a IrSpanned<ExprCompiled>>,
66    bc: &mut BcWriter,
67    k: impl FnOnce(BcSlotInRange, &mut BcWriter),
68) {
69    let exprs: Vec<_> = exprs.into_iter().collect();
70
71    if let Some(slots) = try_slot_range(exprs.iter().copied(), bc) {
72        k(slots, bc);
73    } else {
74        bc.alloc_slots_for_exprs(exprs, |slot, expr, bc| expr.write_bc(slot.to_out(), bc), k)
75    }
76}
77
78pub(crate) fn write_expr_opt(
79    expr: &Option<IrSpanned<ExprCompiled>>,
80    bc: &mut BcWriter,
81    k: impl FnOnce(Option<BcSlotIn>, &mut BcWriter),
82) {
83    if let Some(expr) = expr {
84        expr.write_bc_cb(bc, |slot, bc| k(Some(slot), bc))
85    } else {
86        k(None, bc)
87    }
88}
89
90pub(crate) fn write_n_exprs<const N: usize>(
91    exprs: [&IrSpanned<ExprCompiled>; N],
92    bc: &mut BcWriter,
93    k: impl FnOnce([BcSlotIn; N], &mut BcWriter),
94) {
95    fn help<const N: usize>(
96        mut filled: [BcSlotIn; N],
97        rem_exprs: &[&IrSpanned<ExprCompiled>],
98        bc: &mut BcWriter,
99        k: impl FnOnce([BcSlotIn; N], &mut BcWriter),
100    ) {
101        match rem_exprs.split_first() {
102            Some((first, rem)) => first.write_bc_cb(bc, |first, bc| {
103                filled[N - rem.len() - 1] = first;
104                help(filled, rem, bc, k)
105            }),
106            None => k(filled, bc),
107        }
108    }
109
110    help([BcSlot(98765).to_in(); N], &exprs, bc, k)
111}
112
113impl ExprCompiled {
114    /// Mark variables which are definitely assigned after execution of this expression.
115    ///
116    /// For example, when this expression if executed:
117    ///
118    /// ```python
119    /// t if c else f
120    /// ```
121    ///
122    /// `c` is definitely assigned (because if it is not, then execution fails),
123    /// but we don't know about `t` or `f` because one of them was not executed.
124    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
125        match self {
126            ExprCompiled::Value(_) => {}
127            ExprCompiled::Local(local) => bc.mark_definitely_assigned(*local),
128            ExprCompiled::LocalCaptured(_) => {}
129            ExprCompiled::Module(_) => {}
130            ExprCompiled::Tuple(xs) | ExprCompiled::List(xs) => {
131                for x in xs {
132                    x.mark_definitely_assigned_after(bc);
133                }
134            }
135            ExprCompiled::Dict(xs) => {
136                for (k, v) in xs {
137                    k.mark_definitely_assigned_after(bc);
138                    v.mark_definitely_assigned_after(bc);
139                }
140            }
141            ExprCompiled::Compr(compr) => compr.mark_definitely_assigned_after(bc),
142            ExprCompiled::If(c_t_f) => {
143                let (c, _t, _f) = &**c_t_f;
144                // Condition is executed unconditionally, so we use it to mark definitely assigned.
145                // But we don't know which of the branches will be executed.
146                c.mark_definitely_assigned_after(bc);
147            }
148            ExprCompiled::Slice(l_a_b_c) => {
149                let (l, a, b, c) = &**l_a_b_c;
150                l.mark_definitely_assigned_after(bc);
151                if let Some(a) = a {
152                    a.mark_definitely_assigned_after(bc);
153                }
154                if let Some(b) = b {
155                    b.mark_definitely_assigned_after(bc);
156                }
157                if let Some(c) = c {
158                    c.mark_definitely_assigned_after(bc);
159                }
160            }
161            ExprCompiled::Builtin1(_op, expr) => {
162                expr.mark_definitely_assigned_after(bc);
163            }
164            ExprCompiled::LogicalBinOp(_op, a_b) => {
165                let (a, b) = &**a_b;
166                // `a` is executed unconditionally, but `b` is not,
167                // so we mark only `a` as definitely assigned.
168                a.mark_definitely_assigned_after(bc);
169                let _ = b;
170            }
171            ExprCompiled::Seq(a_b) => {
172                let (a, b) = &**a_b;
173                a.mark_definitely_assigned_after(bc);
174                b.mark_definitely_assigned_after(bc);
175            }
176            ExprCompiled::Builtin2(_op, a_b) => {
177                let (a, b) = &**a_b;
178                a.mark_definitely_assigned_after(bc);
179                b.mark_definitely_assigned_after(bc);
180            }
181            ExprCompiled::Index2(a_i0_i1) => {
182                let (a, i0, i1) = &**a_i0_i1;
183                a.mark_definitely_assigned_after(bc);
184                i0.mark_definitely_assigned_after(bc);
185                i1.mark_definitely_assigned_after(bc);
186            }
187            ExprCompiled::Call(c) => c.mark_definitely_assigned_after(bc),
188            ExprCompiled::Def(d) => d.mark_definitely_assigned_after(bc),
189        }
190    }
191}
192
193impl IrSpanned<ExprCompiled> {
194    fn try_dict_of_consts(
195        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
196    ) -> Option<SmallMap<FrozenValue, FrozenValue>> {
197        let mut res = SmallMap::new();
198        for (k, v) in xs {
199            let k = k.as_value()?.get_hashed().ok()?;
200            let v = v.as_value()?;
201            let prev = res.insert_hashed(k, v);
202            if prev.is_some() {
203                // If there are duplicates, so don't take the fast-literal
204                // path and go down the slow runtime path (which will raise the error).
205                // We have a lint that will likely fire on this issue (and others).
206                return None;
207            }
208        }
209        Some(res)
210    }
211
212    fn try_dict_const_keys(
213        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
214    ) -> Option<Box<[Hashed<FrozenValue>]>> {
215        let mut keys = Vec::new();
216        let mut keys_unique = HashSet::new();
217        for (k, _) in xs {
218            let k = k.as_value()?.get_hashed().ok()?;
219            keys.push(k);
220            let inserted = keys_unique.insert(k);
221            if !inserted {
222                // Otherwise fail at runtime
223                return None;
224            }
225        }
226        Some(keys.into_boxed_slice())
227    }
228
229    fn write_dict(
230        span: FrameSpan,
231        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
232        target: BcSlotOut,
233        bc: &mut BcWriter,
234    ) {
235        if xs.is_empty() {
236            bc.write_instr::<InstrDictNew>(span, target);
237        } else if let Some(d) = Self::try_dict_of_consts(xs) {
238            bc.write_instr::<InstrDictOfConsts>(span, (d, target));
239        } else if let Some(keys) = Self::try_dict_const_keys(xs) {
240            assert_eq!(keys.len(), xs.len());
241            write_exprs(xs.iter().map(|(_, v)| v), bc, |values, bc| {
242                assert_eq!(values.len() as usize, keys.len());
243                bc.write_instr::<InstrDictConstKeys>(span, (keys, values.to_range_from(), target));
244            });
245        } else {
246            let key_spans = xs.map(|(k, _v)| k.span);
247            write_exprs(xs.iter().flat_map(|(k, v)| [k, v]), bc, |kvs, bc| {
248                bc.write_instr_explicit::<InstrDictNPop>(
249                    BcInstrSlowArg {
250                        span,
251                        spans: key_spans,
252                    },
253                    (kvs, target),
254                );
255            });
256        }
257    }
258
259    fn write_not(expr: &IrSpanned<ExprCompiled>, target: BcSlotOut, bc: &mut BcWriter) {
260        expr.write_bc_cb(bc, |slot, bc| {
261            bc.write_instr::<InstrNot>(expr.span, (slot, target));
262        });
263    }
264
265    fn write_equals_const(
266        span: FrameSpan,
267        a: &IrSpanned<ExprCompiled>,
268        b: FrozenValue,
269        target: BcSlotOut,
270        bc: &mut BcWriter,
271    ) {
272        a.write_bc_cb(bc, |a, bc| {
273            if let Some(b) = b.to_value().unpack_int_value() {
274                bc.write_instr::<InstrEqInt>(span, (a, b, target));
275            } else if b.eq_is_ptr_eq() {
276                bc.write_instr::<InstrEqPtr>(span, (a, b, target));
277            } else if let Some(b) = FrozenStringValue::new(b) {
278                bc.write_instr::<InstrEqStr>(span, (a, b, target));
279            } else if let Some(b) = FrozenValueNotSpecial::new(b) {
280                bc.write_instr::<InstrEqConst>(span, (a, b, target));
281            } else {
282                unreachable!("FrozenValue must be either i32, str or not-special");
283            }
284        });
285    }
286
287    fn write_equals(
288        span: FrameSpan,
289        a: &IrSpanned<ExprCompiled>,
290        b: &IrSpanned<ExprCompiled>,
291        target: BcSlotOut,
292        bc: &mut BcWriter,
293    ) {
294        if let Some(a) = a.as_value() {
295            Self::write_equals_const(span, b, a, target, bc);
296        } else if let Some(b) = b.as_value() {
297            Self::write_equals_const(span, a, b, target, bc);
298        } else {
299            write_n_exprs([a, b], bc, |[a, b], bc| {
300                bc.write_instr::<InstrEq>(span, (a, b, target));
301            });
302        }
303    }
304
305    pub(crate) fn write_bc(&self, target: BcSlotOut, bc: &mut BcWriter) {
306        let span = self.span;
307        match &self.node {
308            ExprCompiled::Value(v) => {
309                bc.write_const(span, *v, target);
310            }
311            ExprCompiled::Local(slot) => {
312                bc.write_load_local(span, *slot, target);
313            }
314            ExprCompiled::LocalCaptured(slot) => {
315                bc.write_load_local_captured(span, *slot, target);
316            }
317            ExprCompiled::Module(slot) => {
318                bc.write_instr::<InstrLoadModule>(span, (*slot, target));
319            }
320            ExprCompiled::Tuple(xs) => {
321                write_exprs(xs, bc, |xs, bc| {
322                    bc.write_instr::<InstrTupleNPop>(span, (xs, target));
323                });
324            }
325            ExprCompiled::List(ref xs) => {
326                if xs.is_empty() {
327                    bc.write_instr::<InstrListNew>(span, target);
328                } else if xs.iter().all(|x| x.as_value().is_some()) {
329                    let content = xs.map(|v| v.as_value().unwrap()).into_boxed_slice();
330                    bc.write_instr::<InstrListOfConsts>(span, (content, target));
331                } else {
332                    write_exprs(xs, bc, |xs, bc| {
333                        bc.write_instr::<InstrListNPop>(span, (xs, target));
334                    });
335                }
336            }
337            ExprCompiled::Dict(ref xs) => Self::write_dict(span, xs, target, bc),
338            ExprCompiled::Compr(ref compr) => compr.write_bc(span, target, bc),
339            ExprCompiled::Slice(l_start_stop_step) => {
340                let (l, start, stop, step) = &**l_start_stop_step;
341                l.write_bc_cb(bc, |l, bc| {
342                    write_expr_opt(start, bc, |start, bc| {
343                        write_expr_opt(stop, bc, |stop, bc| {
344                            write_expr_opt(step, bc, |step, bc| {
345                                bc.write_instr::<InstrSlice>(span, (l, start, stop, step, target))
346                            })
347                        })
348                    })
349                });
350            }
351            ExprCompiled::Builtin1(Builtin1::Not, expr) => Self::write_not(expr, target, bc),
352            ExprCompiled::Builtin1(op, expr) => {
353                expr.write_bc_cb(bc, |expr, bc| {
354                    let arg = (expr, target);
355                    match op {
356                        Builtin1::Not => unreachable!("handled above"),
357                        Builtin1::Minus => bc.write_instr::<InstrMinus>(span, arg),
358                        Builtin1::Plus => bc.write_instr::<InstrPlus>(span, arg),
359                        Builtin1::BitNot => bc.write_instr::<InstrBitNot>(span, arg),
360                        Builtin1::TypeIs(t) => {
361                            bc.write_instr::<InstrTypeIs>(span, (expr, *t, target))
362                        }
363                        Builtin1::PercentSOne(before, after) => bc
364                            .write_instr::<InstrPercentSOne>(span, (*before, expr, *after, target)),
365                        Builtin1::FormatOne(before, after) => {
366                            bc.write_instr::<InstrFormatOne>(span, (*before, expr, *after, target))
367                        }
368                        Builtin1::Dot(field) => {
369                            bc.write_instr::<InstrObjectField>(span, (expr, field.clone(), target))
370                        }
371                    }
372                });
373            }
374            ExprCompiled::If(cond_t_f) => {
375                let (cond, t, f) = &**cond_t_f;
376                write_if_else(
377                    cond,
378                    |bc| t.write_bc(target, bc),
379                    |bc| f.write_bc(target, bc),
380                    bc,
381                );
382            }
383            ExprCompiled::LogicalBinOp(op, l_r) => {
384                let (l, r) = &**l_r;
385                l.write_bc_cb(bc, |l_slot, bc| {
386                    let maybe_not = match op {
387                        ExprLogicalBinOp::And => MaybeNot::Id,
388                        ExprLogicalBinOp::Or => MaybeNot::Not,
389                    };
390                    bc.write_if_else(
391                        l_slot,
392                        maybe_not,
393                        l.span,
394                        |bc| r.write_bc(target, bc),
395                        |bc| bc.write_mov(span, l_slot, target),
396                    );
397                });
398            }
399            ExprCompiled::Seq(l_r) => {
400                let (l, r) = &**l_r;
401                l.write_bc_for_effect(bc);
402                r.write_bc(target, bc);
403            }
404            ExprCompiled::Builtin2(Builtin2::Equals, l_r) => {
405                let (l, r) = &**l_r;
406                Self::write_equals(span, l, r, target, bc)
407            }
408            ExprCompiled::Builtin2(op, l_r) => {
409                let (l, r) = &**l_r;
410                write_n_exprs([l, r], bc, |[l, r], bc| {
411                    let arg = (l, r, target);
412                    match op {
413                        Builtin2::Equals => unreachable!("handled above"),
414                        Builtin2::Compare(CompareOp::Less) => {
415                            bc.write_instr::<InstrLess>(span, arg)
416                        }
417                        Builtin2::Compare(CompareOp::Greater) => {
418                            bc.write_instr::<InstrGreater>(span, arg)
419                        }
420                        Builtin2::Compare(CompareOp::LessOrEqual) => {
421                            bc.write_instr::<InstrLessOrEqual>(span, arg)
422                        }
423                        Builtin2::Compare(CompareOp::GreaterOrEqual) => {
424                            bc.write_instr::<InstrGreaterOrEqual>(span, arg)
425                        }
426                        Builtin2::In => bc.write_instr::<InstrIn>(span, arg),
427                        Builtin2::Sub => bc.write_instr::<InstrSub>(span, arg),
428                        Builtin2::Add => bc.write_instr::<InstrAdd>(span, arg),
429                        Builtin2::Multiply => bc.write_instr::<InstrMultiply>(span, arg),
430                        Builtin2::Divide => bc.write_instr::<InstrDivide>(span, arg),
431                        Builtin2::FloorDivide => bc.write_instr::<InstrFloorDivide>(span, arg),
432                        Builtin2::Percent => bc.write_instr::<InstrPercent>(span, arg),
433                        Builtin2::BitAnd => bc.write_instr::<InstrBitAnd>(span, arg),
434                        Builtin2::BitOr => bc.write_instr::<InstrBitOr>(span, arg),
435                        Builtin2::BitXor => bc.write_instr::<InstrBitXor>(span, arg),
436                        Builtin2::LeftShift => bc.write_instr::<InstrLeftShift>(span, arg),
437                        Builtin2::RightShift => bc.write_instr::<InstrRightShift>(span, arg),
438                        Builtin2::ArrayIndex => bc.write_instr::<InstrArrayIndex>(span, arg),
439                    }
440                });
441            }
442            ExprCompiled::Index2(a_i0_i1) => {
443                let (a, i0, i1) = &**a_i0_i1;
444                write_n_exprs([a, i0, i1], bc, |[a, i0, i1], bc| {
445                    bc.write_instr::<InstrArrayIndex2>(span, (a, i0, i1, target))
446                });
447            }
448            ExprCompiled::Call(ref call) => call.write_bc(target, bc),
449            ExprCompiled::Def(ref def) => def.write_bc(span, target, bc),
450        }
451    }
452
453    /// Allocate temporary slot, write expression into it,
454    /// and then consume the slot with the callback.
455    pub(crate) fn write_bc_cb<R>(
456        &self,
457        bc: &mut BcWriter,
458        k: impl FnOnce(BcSlotIn, &mut BcWriter) -> R,
459    ) -> R {
460        if let Some(local) = self.as_local_non_captured() {
461            // Local is known to be definitely assigned, so there's no need
462            // to "load" it just to trigger check that it is assigned.
463            if let Some(slot) = bc.try_definitely_assigned(local) {
464                return k(slot, bc);
465            }
466        }
467
468        bc.alloc_slot(|slot, bc| {
469            self.write_bc(slot.to_out(), bc);
470            k(slot.to_in(), bc)
471        })
472    }
473
474    pub(crate) fn write_bc_for_effect(&self, bc: &mut BcWriter) {
475        self.write_bc_cb(bc, |slot, _bc| {
476            let _ = slot;
477        });
478    }
479}
starlark/eval/bc/compiler/expr.rs

starlark/eval/bc/compiler/
expr.rs
