Implement Patch transpiling

compiler/erg_compiler/codegen.rs

@@ -47,8 +47,8 @@ use AccessKind::*;
 
 /// patch method -> function
 /// patch attr -> variable
-fn debind(name: Option<&str>) -> Option<Str> {
+fn debind(ident: &Identifier) -> Option<Str> {
-    match name {
+    match ident.vi.py_name.as_ref().map(|s| &s[..]) {
         Some(name) if name.starts_with("Function::") => {
             Some(Str::from(name.replace("Function::", "")))
         }
@@ -875,7 +875,7 @@ impl PyCodeGenerator {
                 if is_record {
                     a.ident.dot = Some(DOT);
                 }
-                if let Some(varname) = debind(a.ident.vi.py_name.as_ref().map(|s| &s[..])) {
+                if let Some(varname) = debind(&a.ident) {
                     a.ident.dot = None;
                     a.ident.name = VarName::from_str(varname);
                     self.emit_load_name_instr(a.ident);
@@ -2012,7 +2012,7 @@ impl PyCodeGenerator {
             } else {
                 return self.emit_call_update_310(obj, args);
             }
-        } else if let Some(func_name) = debind(method_name.vi.py_name.as_ref().map(|s| &s[..])) {
+        } else if let Some(func_name) = debind(&method_name) {
             return self.emit_call_fake_method(obj, func_name, method_name, args);
         }
         let is_py_api = obj.is_py_api();

compiler/erg_compiler/transpile.rs

@@ -17,8 +17,8 @@ use crate::context::{Context, ContextProvider};
 use crate::desugar_hir::HIRDesugarer;
 use crate::error::{CompileError, CompileErrors};
 use crate::hir::{
-    Accessor, Array, Block, Call, ClassDef, Def, Dict, Expr, Identifier, Lambda, Params, Set,
+    Accessor, Args, Array, BinOp, Block, Call, ClassDef, Def, Dict, Expr, Identifier, Lambda,
-    Signature, Tuple, HIR,
+    Params, PatchDef, Record, Set, Signature, Tuple, HIR,
 };
 use crate::link::Linker;
 use crate::mod_cache::SharedModuleCache;
@@ -26,6 +26,30 @@ use crate::ty::value::ValueObj;
 use crate::ty::Type;
 use crate::varinfo::VarInfo;
 
+/// patch method -> function
+/// patch attr -> variable
+fn debind(ident: &Identifier) -> Option<Str> {
+    match ident.vi.py_name.as_ref().map(|s| &s[..]) {
+        Some(name) if name.starts_with("Function::") => {
+            Some(Str::from(name.replace("Function::", "")))
+        }
+        Some(patch_method) if patch_method.contains("::") || patch_method.contains('.') => {
+            if ident.vis().is_private() {
+                Some(Str::from(format!("{patch_method}__")))
+            } else {
+                Some(Str::rc(patch_method))
+            }
+        }
+        _ => None,
+    }
+}
+
+fn demangle(name: &str) -> String {
+    name.trim_start_matches("::<module>")
+        .replace("::", "__")
+        .replace('.', "_")
+}
+
 #[derive(Debug)]
 pub enum LastLineOperation {
     Discard,
@@ -267,11 +291,135 @@ impl ScriptGenerator {
         match expr {
             Expr::Lit(lit) => lit.token.content.to_string(),
             Expr::Call(call) => self.transpile_call(call),
-            Expr::BinOp(bin) => match bin.op.kind {
+            Expr::BinOp(bin) => self.transpile_binop(bin),
-                TokenKind::Closed
+            Expr::UnaryOp(unary) => {
-                | TokenKind::LeftOpen
+                let mut code = "(".to_string();
-                | TokenKind::RightOpen
+                if unary.op.kind != TokenKind::Mutate {
-                | TokenKind::Open => {
+                    code += &unary.op.content;
+                }
+                code += &self.transpile_expr(*unary.expr);
+                code += ")";
+                code
+            }
+            Expr::Array(array) => match array {
+                Array::Normal(arr) => {
+                    let mut code = "[".to_string();
+                    for elem in arr.elems.pos_args {
+                        code += &format!("{},", self.transpile_expr(elem.expr));
+                    }
+                    code += "]";
+                    code
+                }
+                other => todo!("transpiling {other}"),
+            },
+            Expr::Set(set) => match set {
+                Set::Normal(st) => {
+                    let mut code = "{".to_string();
+                    for elem in st.elems.pos_args {
+                        code += &format!("{},", self.transpile_expr(elem.expr));
+                    }
+                    code += "}";
+                    code
+                }
+                other => todo!("transpiling {other}"),
+            },
+            Expr::Record(rec) => self.transpile_record(rec),
+            Expr::Tuple(tuple) => match tuple {
+                Tuple::Normal(tup) => {
+                    let mut code = "(".to_string();
+                    for elem in tup.elems.pos_args {
+                        code += &format!("{},", self.transpile_expr(elem.expr));
+                    }
+                    code += ")";
+                    code
+                }
+            },
+            Expr::Dict(dict) => match dict {
+                Dict::Normal(dic) => {
+                    let mut code = "{".to_string();
+                    for kv in dic.kvs {
+                        code += &format!(
+                            "({}): ({}),",
+                            self.transpile_expr(kv.key),
+                            self.transpile_expr(kv.value)
+                        );
+                    }
+                    code += "}";
+                    code
+                }
+                other => todo!("transpiling {other}"),
+            },
+            Expr::Accessor(acc) => self.transpile_acc(acc),
+            Expr::Def(def) => self.transpile_def(def),
+            Expr::Lambda(lambda) => self.transpile_lambda(lambda),
+            Expr::ClassDef(classdef) => self.transpile_classdef(classdef),
+            Expr::PatchDef(patchdef) => self.transpile_patchdef(patchdef),
+            Expr::AttrDef(mut adef) => {
+                let mut code = format!("{} = ", self.transpile_expr(Expr::Accessor(adef.attr)));
+                if adef.block.len() > 1 {
+                    let name = format!("instant_block_{}__", self.fresh_var_n);
+                    self.fresh_var_n += 1;
+                    let mut code = format!("def {name}():\n");
+                    code += &self.transpile_block(adef.block, Return);
+                    self.prelude += &code;
+                    format!("{name}()")
+                } else {
+                    let expr = adef.block.remove(0);
+                    code += &self.transpile_expr(expr);
+                    code
+                }
+            }
+            // TODO:
+            Expr::Compound(comp) => {
+                let mut code = "".to_string();
+                for expr in comp.into_iter() {
+                    code += &self.transpile_expr(expr);
+                    code += &format!("\n{}", "    ".repeat(self.level));
+                }
+                code
+            }
+            Expr::Import(acc) => {
+                let full_name = Str::from(acc.show());
+                let root = PyCodeGenerator::get_root(&acc);
+                self.prelude += &format!(
+                    "{} = __import__(\"{full_name}\")\n",
+                    Self::transpile_ident(root)
+                );
+                String::new()
+            }
+            other => todo!("transpile {other}"),
+        }
+    }
+
+    fn transpile_record(&mut self, rec: Record) -> String {
+        if !self.namedtuple_loaded {
+            self.load_namedtuple();
+            self.namedtuple_loaded = true;
+        }
+        let mut attrs = "[".to_string();
+        let mut values = "(".to_string();
+        for mut attr in rec.attrs.into_iter() {
+            attrs += &format!("'{}',", Self::transpile_ident(attr.sig.into_ident()));
+            if attr.body.block.len() > 1 {
+                let name = format!("instant_block_{}__", self.fresh_var_n);
+                self.fresh_var_n += 1;
+                let mut code = format!("def {name}():\n");
+                code += &self.transpile_block(attr.body.block, Return);
+                self.prelude += &code;
+                values += &format!("{name}(),");
+            } else {
+                let expr = attr.body.block.remove(0);
+                values += &format!("{},", self.transpile_expr(expr));
+            }
+        }
+        attrs += "]";
+        values += ")";
+        format!("NamedTuple__('Record', {attrs}){values}")
+    }
+
+    fn transpile_binop(&mut self, bin: BinOp) -> String {
+        match bin.op.kind {
+            TokenKind::Closed | TokenKind::LeftOpen | TokenKind::RightOpen | TokenKind::Open => {
                 if !self.range_ops_loaded {
                     self.load_range_ops();
                     self.range_ops_loaded = true;
@@ -312,89 +460,11 @@ impl ScriptGenerator {
                 code += ")";
                 code
             }
-            },
-            Expr::UnaryOp(unary) => {
-                let mut code = "(".to_string();
-                if unary.op.kind != TokenKind::Mutate {
-                    code += &unary.op.content;
-                }
-                code += &self.transpile_expr(*unary.expr);
-                code += ")";
-                code
-            }
-            Expr::Array(array) => match array {
-                Array::Normal(arr) => {
-                    let mut code = "[".to_string();
-                    for elem in arr.elems.pos_args {
-                        code += &format!("{},", self.transpile_expr(elem.expr));
-                    }
-                    code += "]";
-                    code
-                }
-                other => todo!("transpiling {other}"),
-            },
-            Expr::Set(set) => match set {
-                Set::Normal(st) => {
-                    let mut code = "{".to_string();
-                    for elem in st.elems.pos_args {
-                        code += &format!("{},", self.transpile_expr(elem.expr));
-                    }
-                    code += "}";
-                    code
-                }
-                other => todo!("transpiling {other}"),
-            },
-            Expr::Record(rec) => {
-                if !self.namedtuple_loaded {
-                    self.load_namedtuple();
-                    self.namedtuple_loaded = true;
-                }
-                let mut attrs = "[".to_string();
-                let mut values = "(".to_string();
-                for mut attr in rec.attrs.into_iter() {
-                    attrs += &format!("'{}',", Self::transpile_ident(attr.sig.into_ident()));
-                    if attr.body.block.len() > 1 {
-                        let name = format!("instant_block_{}__", self.fresh_var_n);
-                        self.fresh_var_n += 1;
-                        let mut code = format!("def {name}():\n");
-                        code += &self.transpile_block(attr.body.block, Return);
-                        self.prelude += &code;
-                        values += &format!("{name}(),");
-                    } else {
-                        let expr = attr.body.block.remove(0);
-                        values += &format!("{},", self.transpile_expr(expr));
         }
     }
-                attrs += "]";
+
-                values += ")";
+    fn transpile_acc(&mut self, acc: Accessor) -> String {
-                format!("NamedTuple__('Record', {attrs}){values}")
+        match acc {
-            }
-            Expr::Tuple(tuple) => match tuple {
-                Tuple::Normal(tup) => {
-                    let mut code = "(".to_string();
-                    for elem in tup.elems.pos_args {
-                        code += &format!("{},", self.transpile_expr(elem.expr));
-                    }
-                    code += ")";
-                    code
-                }
-            },
-            Expr::Dict(dict) => match dict {
-                Dict::Normal(dic) => {
-                    let mut code = "{".to_string();
-                    for kv in dic.kvs {
-                        code += &format!(
-                            "({}): ({}),",
-                            self.transpile_expr(kv.key),
-                            self.transpile_expr(kv.value)
-                        );
-                    }
-                    code += "}";
-                    code
-                }
-                other => todo!("transpiling {other}"),
-            },
-            Expr::Accessor(acc) => match acc {
             Accessor::Ident(ident) => {
                 match &ident.inspect()[..] {
                     "Str" | "Bool" | "Nat" | "Array" if !self.builtin_types_loaded => {
@@ -406,51 +476,17 @@ impl ScriptGenerator {
                 Self::transpile_ident(ident)
             }
             Accessor::Attr(attr) => {
+                if let Some(name) = debind(&attr.ident) {
+                    demangle(&name)
+                } else {
                     format!(
                         "({}).{}",
                         self.transpile_expr(*attr.obj),
                         Self::transpile_ident(attr.ident)
                     )
                 }
-            },
-            Expr::Def(def) => self.transpile_def(def),
-            Expr::Lambda(lambda) => self.transpile_lambda(lambda),
-            Expr::ClassDef(classdef) => self.transpile_classdef(classdef),
-            Expr::AttrDef(mut adef) => {
-                let mut code = format!("{} = ", self.transpile_expr(Expr::Accessor(adef.attr)));
-                if adef.block.len() > 1 {
-                    let name = format!("instant_block_{}__", self.fresh_var_n);
-                    self.fresh_var_n += 1;
-                    let mut code = format!("def {name}():\n");
-                    code += &self.transpile_block(adef.block, Return);
-                    self.prelude += &code;
-                    format!("{name}()")
-                } else {
-                    let expr = adef.block.remove(0);
-                    code += &self.transpile_expr(expr);
-                    code
             }
         }
-            // TODO:
-            Expr::Compound(comp) => {
-                let mut code = "".to_string();
-                for expr in comp.into_iter() {
-                    code += &self.transpile_expr(expr);
-                    code += &format!("\n{}", "    ".repeat(self.level));
-                }
-                code
-            }
-            Expr::Import(acc) => {
-                let full_name = Str::from(acc.show());
-                let root = PyCodeGenerator::get_root(&acc);
-                self.prelude += &format!(
-                    "{} = __import__(\"{full_name}\")\n",
-                    Self::transpile_ident(root)
-                );
-                String::new()
-            }
-            other => todo!("transpile {other:?}"),
-        }
     }
 
     fn transpile_call(&mut self, mut call: Call) -> String {
@@ -462,7 +498,33 @@ impl ScriptGenerator {
                 }
                 code
             }
-            Some("if" | "if!") => {
+            Some("not") => format!("(not ({}))", self.transpile_expr(call.args.remove(0))),
+            Some("if" | "if!") => self.transpile_if(call),
+            Some("for" | "for!") => {
+                let mut code = "for ".to_string();
+                let iter = call.args.remove(0);
+                let Expr::Lambda(block) = call.args.remove(0) else { todo!() };
+                let sig = block.params.non_defaults.get(0).unwrap();
+                let ParamPattern::VarName(param) = &sig.pat else { todo!() };
+                code += &format!("{}__ ", &param.token().content);
+                code += &format!("in {}:\n", self.transpile_expr(iter));
+                code += &self.transpile_block(block.body, Discard);
+                code
+            }
+            Some("while" | "while!") => {
+                let mut code = "while ".to_string();
+                let cond = call.args.remove(0);
+                let Expr::Lambda(block) = call.args.remove(0) else { todo!() };
+                code += &format!("{}:\n", self.transpile_expr(cond));
+                code += &self.transpile_block(block.body, Discard);
+                code
+            }
+            Some("match" | "match!") => self.transpile_match(call),
+            _ => self.transpile_simple_call(call),
+        }
+    }
+
+    fn transpile_if(&mut self, mut call: Call) -> String {
         let cond = self.transpile_expr(call.args.remove(0));
         let Expr::Lambda(mut then_block) = call.args.remove(0) else { todo!() };
         let else_block = call.args.try_remove(0).map(|ex| {
@@ -508,26 +570,8 @@ impl ScriptGenerator {
         // This is a very bad design, but can be used for this code
         format!("{tmp_func}()")
     }
-            Some("for" | "for!") => {
+
-                let mut code = "for ".to_string();
+    fn transpile_match(&mut self, mut call: Call) -> String {
-                let iter = call.args.remove(0);
-                let Expr::Lambda(block) = call.args.remove(0) else { todo!() };
-                let sig = block.params.non_defaults.get(0).unwrap();
-                let ParamPattern::VarName(param) = &sig.pat else { todo!() };
-                code += &format!("{}__ ", &param.token().content);
-                code += &format!("in {}:\n", self.transpile_expr(iter));
-                code += &self.transpile_block(block.body, Discard);
-                code
-            }
-            Some("while" | "while!") => {
-                let mut code = "while ".to_string();
-                let cond = call.args.remove(0);
-                let Expr::Lambda(block) = call.args.remove(0) else { todo!() };
-                code += &format!("{}:\n", self.transpile_expr(cond));
-                code += &self.transpile_block(block.body, Discard);
-                code
-            }
-            Some("match" | "match!") => {
         let tmp = Str::from(format!("match_tmp_{}__", self.fresh_var_n));
         self.fresh_var_n += 1;
         let tmp_func = Str::from(format!("match_tmp_func_{}__", self.fresh_var_n));
@@ -574,13 +618,19 @@ impl ScriptGenerator {
         self.prelude += &code;
         format!("{tmp_func}()")
     }
-            _ => self.transpile_simple_call(call),
-        }
-    }
 
-    fn transpile_simple_call(&mut self, mut call: Call) -> String {
+    fn transpile_simple_call(&mut self, call: Call) -> String {
         let is_py_api = if let Some(attr) = &call.attr_name {
-            attr.is_py_api()
+            let is_py_api = attr.is_py_api();
+            if let Some(name) = debind(attr) {
+                let name = demangle(&name);
+                return format!(
+                    "{name}({}, {})",
+                    self.transpile_expr(*call.obj),
+                    self.transpile_args(call.args, is_py_api, false)
+                );
+            }
+            is_py_api
         } else {
             call.obj.is_py_api()
         };
@@ -588,12 +638,20 @@ impl ScriptGenerator {
         if let Some(attr) = call.attr_name {
             code += &format!(".{}", Self::transpile_ident(attr));
         }
+        code += &self.transpile_args(call.args, is_py_api, true);
+        code
+    }
+
+    fn transpile_args(&mut self, mut args: Args, is_py_api: bool, paren: bool) -> String {
+        let mut code = String::new();
+        if paren {
             code.push('(');
-        while let Some(arg) = call.args.try_remove_pos(0) {
+        }
+        while let Some(arg) = args.try_remove_pos(0) {
             code += &self.transpile_expr(arg.expr);
             code.push(',');
         }
-        while let Some(arg) = call.args.try_remove_kw(0) {
+        while let Some(arg) = args.try_remove_kw(0) {
             let escape = if is_py_api { "" } else { "__" };
             code += &format!(
                 "{}{escape}={},",
@@ -601,13 +659,15 @@ impl ScriptGenerator {
                 self.transpile_expr(arg.expr)
             );
         }
+        if paren {
             code.push(')');
+        }
         code
     }
 
     fn transpile_ident(ident: Identifier) -> String {
         if let Some(py_name) = ident.vi.py_name {
-            py_name.to_string()
+            demangle(&py_name)
         } else if ident.dot.is_some() {
             ident.name.into_token().content.to_string()
         } else {
@@ -732,4 +792,21 @@ impl ScriptGenerator {
         code += &self.transpile_block(classdef.methods, Discard);
         code
     }
+
+    fn transpile_patchdef(&mut self, patch_def: PatchDef) -> String {
+        let mut code = String::new();
+        for chunk in patch_def.methods.into_iter() {
+            let Expr::Def(mut def) = chunk else { todo!() };
+            let name = format!(
+                "{}{}",
+                demangle(&patch_def.sig.ident().to_string_without_type()),
+                demangle(&def.sig.ident().to_string_without_type()),
+            );
+            def.sig.ident_mut().name = VarName::from_str(Str::from(name));
+            code += &"    ".repeat(self.level);
+            code += &self.transpile_def(def);
+            code.push('\n');
+        }
+        code
+    }
 }