<commit_before>//! Various inlining transforms.
//!
//! These transforms take a set of nested expressions and fuse them into a single one.

extern crate fnv;

use ast::*;
use ast::ExprKind::*;

use fnv::FnvHashMap;

use std::mem;

#[cfg(test)]
use tests::*;

/// Inlines GetField(MakeStruct(*)) expressions, which can occur during loop fusion when some
/// of the loops are zipping together multiple column vectors.
pub fn inline_get_field(expr: &mut Expr) {
    expr.transform(&mut |ref mut expr| {
        if let GetField { ref expr, index } = expr.kind {
            if let MakeStruct { ref elems } = expr.kind {
                return Some(elems[index as usize].clone());
            }
        }
        None
    });
}

/// Inlines Zip expressions as collections of iters. Using Zips outside of a For loop is currently
/// unsupported behavior. This transform handles the simple case of converting Zips in macros
/// such as map and filter into Iters in For loops.
///
/// TODO(shoumik): Perhaps Zip should just be a macro? Then macros need to be ordered.
pub fn inline_zips(expr: &mut Expr) {
    expr.transform(&mut |ref mut e| {
        if let For {
                   ref mut iters,
                   ref builder,
                   ref func,
               } = e.kind {
            if iters.len() == 1 {
                let ref first_iter = iters[0];
                if let Zip { ref vectors } = first_iter.data.kind {
                    let new_iters = vectors
                        .iter()
                        .map(|v| {
                                 Iter {
                                     data: Box::new(v.clone()),
                                     start: None,
                                     end: None,
                                     stride: None,
                                     kind: first_iter.kind.clone(),
                                     shape: None,
                                     strides: None,
                                 }
                             })
                        .collect::<Vec<_>>();
                    return Some(Expr {
                                    ty: e.ty.clone(),
                                    kind: For {
                                        iters: new_iters,
                                        builder: builder.clone(),
                                        func: func.clone(),
                                    },
                                    annotations: Annotations::new(),
                                });
                }
            }
        }
        None
    });
}

/// Inlines Apply nodes whose argument is a Lambda expression. These often arise during macro
/// expansion but it's simpler to inline them before doing type inference.
/// Unlike many of the other transformations, we make this one independent of types so that
/// we can apply it before type inference.
///
/// Caveats:
/// - Functions that reuse a parameter twice have its expansion appear twice, instead of assigning
///   it to a temporary as would happen with function application.
/// - Does not complete inlining if some of the functions take functions as arguments (in that
///   case, the expressions after inlining may lead to more inlining).
pub fn inline_apply(expr: &mut Expr) {
    expr.transform(&mut |ref mut expr| {
        if let Apply {
                   ref func,
                   params: ref args,
               } = expr.kind {
            if let Lambda {
                       ref params,
                       ref body,
                   } = func.kind {
                let mut new = *body.clone();
                for (param, arg) in params.iter().zip(args) {
                    new.substitute(&param.name, &arg);
                }
                return Some(new);
            }
        }
        None
    });
}

pub fn inline_let(expr: &mut Expr) {
    expr.uniquify().unwrap();
    let ref mut usages = FnvHashMap::default();
    count_symbols(expr, usages);
    trace!("Symbol count: {:?}", usages);
    inline_let_helper(expr, usages)
}

#[derive(Debug)]
struct SymbolTracker {
    count: i32,
    loop_nest: i32,
    value: Option<Box<Expr>>,
}

impl Default for SymbolTracker {
    fn default() -> SymbolTracker {
        SymbolTracker {
            count: 0,
            loop_nest: 0,
            value: None
        }
    }
}

/// Count the occurances of each symbol defined by a `Let` statement.
fn count_symbols(expr: &Expr, usage: &mut FnvHashMap<Symbol, SymbolTracker>) {
    match expr.kind {
        For { ref func, .. } | Iterate { update_func: ref func, .. } => {
            // Mark all symbols seen so far as "in a loop"
            for value in usage.values_mut() {
                value.loop_nest += 1;
            }

            count_symbols(func, usage);

            for value in usage.values_mut() {
                value.loop_nest -= 1;
            }
        }
        Let { ref name, .. } => {
            debug_assert!(!usage.contains_key(name));
            let _ = usage.insert(name.clone(), SymbolTracker::default());
        }
        Ident(ref symbol) => {
            if let Some(ref mut tracker) = usage.get_mut(symbol) {<commit_msg>Sort fix for inliner (#412)<commit_after>//! Various inlining transforms.
//!
//! These transforms take a set of nested expressions and fuse them into a single one.

extern crate fnv;

use ast::*;
use ast::ExprKind::*;

use fnv::FnvHashMap;

use std::mem;

#[cfg(test)]
use tests::*;

/// Inlines GetField(MakeStruct(*)) expressions, which can occur during loop fusion when some
/// of the loops are zipping together multiple column vectors.
pub fn inline_get_field(expr: &mut Expr) {
    expr.transform(&mut |ref mut expr| {
        if let GetField { ref expr, index } = expr.kind {
            if let MakeStruct { ref elems } = expr.kind {
                return Some(elems[index as usize].clone());
            }
        }
        None
    });
}

/// Inlines Zip expressions as collections of iters. Using Zips outside of a For loop is currently
/// unsupported behavior. This transform handles the simple case of converting Zips in macros
/// such as map and filter into Iters in For loops.
///
/// TODO(shoumik): Perhaps Zip should just be a macro? Then macros need to be ordered.
pub fn inline_zips(expr: &mut Expr) {
    expr.transform(&mut |ref mut e| {
        if let For {
                   ref mut iters,
                   ref builder,
                   ref func,
               } = e.kind {
            if iters.len() == 1 {
                let ref first_iter = iters[0];
                if let Zip { ref vectors } = first_iter.data.kind {
                    let new_iters = vectors
                        .iter()
                        .map(|v| {
                                 Iter {
                                     data: Box::new(v.clone()),
                                     start: None,
                                     end: None,
                                     stride: None,
                                     kind: first_iter.kind.clone(),
                                     shape: None,
                                     strides: None,
                                 }
                             })
                        .collect::<Vec<_>>();
                    return Some(Expr {
                                    ty: e.ty.clone(),
                                    kind: For {
                                        iters: new_iters,
                                        builder: builder.clone(),
                                        func: func.clone(),
                                    },
                                    annotations: Annotations::new(),
                                });
                }
            }
        }
        None
    });
}

/// Inlines Apply nodes whose argument is a Lambda expression. These often arise during macro
/// expansion but it's simpler to inline them before doing type inference.
/// Unlike many of the other transformations, we make this one independent of types so that
/// we can apply it before type inference.
///
/// Caveats:
/// - Functions that reuse a parameter twice have its expansion appear twice, instead of assigning
///   it to a temporary as would happen with function application.
/// - Does not complete inlining if some of the functions take functions as arguments (in that
///   case, the expressions after inlining may lead to more inlining).
pub fn inline_apply(expr: &mut Expr) {
    expr.transform(&mut |ref mut expr| {
        if let Apply {
                   ref func,
                   params: ref args,
               } = expr.kind {
            if let Lambda {
                       ref params,
                       ref body,
                   } = func.kind {
                let mut new = *body.clone();
                for (param, arg) in params.iter().zip(args) {
                    new.substitute(&param.name, &arg);
                }
                return Some(new);
            }
        }
        None
    });
}

pub fn inline_let(expr: &mut Expr) {
    expr.uniquify().unwrap();
    let ref mut usages = FnvHashMap::default();
    count_symbols(expr, usages);
    trace!("Symbol count: {:?}", usages);
    inline_let_helper(expr, usages)
}

#[derive(Debug)]
struct SymbolTracker {
    count: i32,
    loop_nest: i32,
    value: Option<Box<Expr>>,
}

impl Default for SymbolTracker {
    fn default() -> SymbolTracker {
        SymbolTracker {
            count: 0,
            loop_nest: 0,
            value: None
        }
    }
}

/// Count the occurances of each symbol defined by a `Let` statement.
fn count_symbols(expr: &Expr, usage: &mut FnvHashMap<Symbol, SymbolTracker>) {
    match expr.kind {
        For { ref func, .. } | Iterate { update_func: ref func, .. } | Sort { keyfunc: ref func, .. } => {
            // Mark all symbols seen so far as "in a loop"
            for value in usage.values_mut() {
                value.loop_nest += 1;
            }

            count_symbols(func, usage);

            for value in usage.values_mut() {
                value.loop_nest -= 1;
            }
        }
        Let { ref name, .. } => {
            debug_assert!(!usage.contains_key(name));
            let _ = usage.insert(name.clone(), SymbolTracker::default());
        }
        Ident(ref symbol) => {
            if let Some(ref mut tracker) = usage.get_mut(symbol) {
{'fraction_non_alphanumeric': 0.09874892560404928, 'fraction_numerical': 0.0021965428325852352, 'mean_word_length': 1.3152774707052841, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 32, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3502460', 'n_tokens_mistral': 2717, 'n_tokens_neox': 2569, 'n_words': 900}