Implement name mangling

src/compile.rs

@@ -5,16 +5,20 @@
 //!   https://github.com/python/cpython/blob/master/Python/compile.c
 //!   https://github.com/micropython/micropython/blob/master/py/compile.c
 
-use crate::error::{CompileError, CompileErrorType};
 use crate::ir::{self, CodeInfo};
 pub use crate::mode::Mode;
 use crate::symboltable::{make_symbol_table, make_symbol_table_expr, SymbolScope, SymbolTable};
 use crate::IndexSet;
+use crate::{
+    error::{CompileError, CompileErrorType},
+    symboltable,
+};
 use itertools::Itertools;
 use num_complex::Complex64;
 use num_traits::ToPrimitive;
 use rustpython_ast as ast;
 use rustpython_bytecode::{self as bytecode, CodeObject, ConstantData, Instruction};
+use std::borrow::Cow;
 
 type CompileResult<T> = Result<T, CompileError>;
 
@@ -33,6 +37,7 @@ struct Compiler {
     current_qualified_path: Option<String>,
     done_with_future_stmts: bool,
     ctx: CompileContext,
+    class_name: Option<String>,
     opts: CompileOpts,
 }
 
@@ -170,6 +175,7 @@ impl Compiler {
                 in_class: false,
                 func: FunctionContext::NoFunction,
             },
+            class_name: None,
             opts,
         }
     }
@@ -260,10 +266,11 @@ impl Compiler {
         name: &str,
         cache: impl FnOnce(&mut CodeInfo) -> &mut IndexSet<String>,
     ) -> bytecode::NameIdx {
+        let name = self.mangle(name);
         let cache = cache(self.current_codeinfo());
         cache
-            .get_index_of(name)
+            .get_index_of(name.as_ref())
-            .unwrap_or_else(|| cache.insert_full(name.to_owned()).0) as u32
+            .unwrap_or_else(|| cache.insert_full(name.into_owned()).0) as u32
     }
 
     fn compile_program(
@@ -358,9 +365,14 @@ impl Compiler {
         self.compile_name(name, NameUsage::Store)
     }
 
+    fn mangle<'a>(&self, name: &'a str) -> Cow<'a, str> {
+        symboltable::mangle_name(self.class_name.as_deref(), name)
+    }
+
     fn compile_name(&mut self, name: &str, usage: NameUsage) {
+        let name = self.mangle(name);
         let symbol_table = self.symbol_table_stack.last().unwrap();
-        let symbol = symbol_table.lookup(name).expect(
+        let symbol = symbol_table.lookup(name.as_ref()).expect(
             "The symbol must be present in the symbol table, even when it is undefined in python.",
         );
         let info = self.code_stack.last_mut().unwrap();
@@ -394,8 +406,8 @@ impl Compiler {
             // SymbolScope::Unknown => NameOpType::Global,
         };
         let mut idx = cache
-            .get_index_of(name)
+            .get_index_of(name.as_ref())
-            .unwrap_or_else(|| cache.insert_full(name.to_owned()).0);
+            .unwrap_or_else(|| cache.insert_full(name.into_owned()).0);
         if let SymbolScope::Free = symbol.scope {
             idx += info.cellvar_cache.len();
         }
@@ -1033,7 +1045,7 @@ impl Compiler {
         for arg in args_iter {
             if let Some(annotation) = &arg.node.annotation {
                 self.emit_constant(ConstantData::Str {
-                    value: arg.node.arg.to_owned(),
+                    value: self.mangle(&arg.node.arg).into_owned(),
                 });
                 self.compile_expression(&annotation)?;
                 num_annotations += 1;
@@ -1153,9 +1165,13 @@ impl Compiler {
             loop_data: None,
         };
 
+        let prev_class_name = std::mem::replace(&mut self.class_name, Some(name.to_owned()));
+
         let qualified_name = self.create_qualified_name(name, "");
-        let old_qualified_path = self.current_qualified_path.take();
+        let old_qualified_path = std::mem::replace(
-        self.current_qualified_path = Some(qualified_name.clone());
+            &mut self.current_qualified_path,
+            Some(qualified_name.clone()),
+        );
 
         self.emit(Instruction::LoadBuildClass);
         self.push_output(bytecode::CodeFlags::empty(), 0, 0, 0, name.to_owned());
@@ -1201,6 +1217,7 @@ impl Compiler {
 
         let code = self.pop_code_object();
 
+        self.class_name = prev_class_name;
         self.current_qualified_path = old_qualified_path;
         self.ctx = prev_ctx;
 
@@ -1531,7 +1548,7 @@ impl Compiler {
             let annotations = self.name("__annotations__");
             self.emit(Instruction::LoadNameAny(annotations));
             self.emit_constant(ConstantData::Str {
-                value: id.to_owned(),
+                value: self.mangle(id).into_owned(),
             });
             self.emit(Instruction::StoreSubscript);
         } else {

src/symboltable.rs

@@ -10,6 +10,7 @@ Inspirational file: https://github.com/python/cpython/blob/master/Python/symtabl
 use crate::error::{CompileError, CompileErrorType};
 use crate::IndexMap;
 use rustpython_ast::{self as ast, Location};
+use std::borrow::Cow;
 use std::fmt;
 
 pub fn make_symbol_table(program: &[ast::Stmt]) -> SymbolTableResult<SymbolTable> {
@@ -533,6 +534,7 @@ enum SymbolUsage {
 }
 
 struct SymbolTableBuilder {
+    class_name: Option<String>,
     // Scope stack.
     tables: Vec<SymbolTable>,
 }
@@ -552,7 +554,10 @@ enum ExpressionContext {
 
 impl SymbolTableBuilder {
     fn new() -> Self {
-        let mut this = Self { tables: vec![] };
+        let mut this = Self {
+            class_name: None,
+            tables: vec![],
+        };
         this.enter_scope("top", SymbolTableType::Module, 0);
         this
     }
@@ -664,12 +669,14 @@ impl SymbolTableBuilder {
                 decorator_list,
             } => {
                 self.enter_scope(name, SymbolTableType::Class, location.row());
+                let prev_class = std::mem::replace(&mut self.class_name, Some(name.to_owned()));
                 self.register_name("__module__", SymbolUsage::Assigned, location)?;
                 self.register_name("__qualname__", SymbolUsage::Assigned, location)?;
                 self.register_name("__doc__", SymbolUsage::Assigned, location)?;
                 self.register_name("__class__", SymbolUsage::Assigned, location)?;
                 self.scan_statements(body)?;
                 self.leave_scope();
+                self.class_name = prev_class;
                 self.scan_expressions(bases, ExpressionContext::Load)?;
                 for keyword in keywords {
                     self.scan_expression(&keyword.node.value, ExpressionContext::Load)?;
@@ -1108,15 +1115,14 @@ impl SymbolTableBuilder {
         let scope_depth = self.tables.len();
         let table = self.tables.last_mut().unwrap();
 
+        let name = mangle_name(self.class_name.as_deref(), name);
+
         // Some checks for the symbol that present on this scope level:
-        let symbol = if let Some(symbol) = table.symbols.get_mut(name) {
+        let symbol = if let Some(symbol) = table.symbols.get_mut(name.as_ref()) {
             // Role already set..
             match role {
                 SymbolUsage::Global => {
-                    if symbol.is_global() {
+                    if !symbol.is_global() {
-                        // Ok
-                        symbol.scope = SymbolScope::GlobalExplicit;
-                    } else {
                         return Err(SymbolTableError {
                             error: format!("name '{}' is used prior to global declaration", name),
                             location,
@@ -1172,8 +1178,8 @@ impl SymbolTableBuilder {
                 }
             }
             // Insert symbol when required:
-            let symbol = Symbol::new(name);
+            let symbol = Symbol::new(name.as_ref());
-            table.symbols.entry(name.to_owned()).or_insert(symbol)
+            table.symbols.entry(name.into_owned()).or_insert(symbol)
         };
 
         // Set proper flags on symbol:
@@ -1205,16 +1211,7 @@ impl SymbolTableBuilder {
                 symbol.is_assign_namedexpr_in_comprehension = true;
             }
             SymbolUsage::Global => {
-                if let SymbolScope::Unknown = symbol.scope {
+                symbol.scope = SymbolScope::GlobalExplicit;
-                    symbol.scope = SymbolScope::GlobalExplicit;
-                } else if symbol.is_global() {
-                    symbol.scope = SymbolScope::GlobalExplicit;
-                } else {
-                    return Err(SymbolTableError {
-                        error: format!("Symbol {} scope cannot be set to global, since its scope was already determined otherwise.", name),
-                        location,
-                    });
-                }
             }
             SymbolUsage::Used => {
                 symbol.is_referenced = true;
@@ -1239,3 +1236,20 @@ impl SymbolTableBuilder {
         Ok(())
     }
 }
+
+pub(crate) fn mangle_name<'a>(class_name: Option<&str>, name: &'a str) -> Cow<'a, str> {
+    let class_name = match class_name {
+        Some(n) => n,
+        None => return name.into(),
+    };
+    if !name.starts_with("__") || name.ends_with("__") || name.contains('.') {
+        return name.into();
+    }
+    // strip leading underscore
+    let class_name = class_name.strip_prefix(|c| c == '_').unwrap_or(class_name);
+    let mut ret = String::with_capacity(1 + class_name.len() + name.len());
+    ret.push('_');
+    ret.push_str(class_name);
+    ret.push_str(name);
+    ret.into()
+}