Skip to content

在 Rust 中使用过程宏模拟 Python 缩进语法

预期形式

python! {
    fn f(a: u32, b: u32) -> u32:
        return a + b
}

python! {
    fn main():
        for ith in 0..7:
            for jth in 0..ith:
                let ans = f(ith, jth)
                println!("{}", ans)
}

fn f(a: u32, b: u32) -> u32 {
    return a + b
}

fn main() {
    for ith in 0..7 {
        for jth in 0..ith {
            let ans = f(ith, jth)
            println!("{}", ans)
        }
    }
}

原理

使用 Rust 中的过程宏，对源代码进行预处理。首先将行末的冒号替换为左大括号，其次通过将缩进数量压栈的方式，计算某一行代码是否属于当前的逻辑范围，进而补全右大括号，从而模拟 Python 的缩进语法，代码如下：

lib.rs 
#![feature(proc_macro_span)]

extern crate proc_macro;
use proc_macro::TokenStream;

use std::cmp::Ordering::{Greater, Less, Equal};

#[proc_macro]
pub fn python(tokens: TokenStream) -> TokenStream {
    let mut lines: Vec<String> = source_code(tokens)
        .lines()
        .map(|line| String::from(line.trim_end()))
        .filter(|line| !line.is_empty())
        .collect();
    let mut stack = vec![indentation_size(&lines[0])];
    for ith in 0..lines.len() {
        // :
        if lines[ith].ends_with(':') {
            lines[ith].pop();
            (lines[ith].push(' '), lines[ith].push('{'));
        } else {
            lines[ith].push(';');
        }
        // }
        let indent = indentation_size(&lines[ith]);
        loop {
            match stack.iter().last() {
                Some(previous) => {
                    match indent.cmp(previous) {
                        Greater => stack.push(indent),
                        Less => {
                            lines[ith - 1].push('}');
                            stack.pop();
                        },
                        Equal => break,
                    };
                }
                None => unreachable!(),
            };
        }
    }
    let suffix: String = (0..stack.len() - 1).map(|_| "}").collect();
    return format!("{}{}", lines.join("\n"), suffix).parse().unwrap();
}

fn source_code(tokens: TokenStream) -> String {
    let mut ans = String::new();
    let (mut line, mut column) = (0, 0);
    for token in tokens {
        let location = token.span().start();
        let string = token.to_string();
        // vertical
        if line < location.line {
            ans.push('\n');
            line = location.line;
            column = 0;
        }
        // horizontal
        while column < location.column {
            ans.push(' ');
            column += 1;
        }
        // token
        ans.push_str(&string);
        column += string.len();
    }
    return ans;
}

fn indentation_size(string: &String) -> usize {
    return string.len() - string.trim_start().len();
}

参考文献