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,52 +291,7 @@ 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
-                | TokenKind::LeftOpen
-                | TokenKind::RightOpen
-                | TokenKind::Open => {
-                    if !self.range_ops_loaded {
-                        self.load_range_ops();
-                        self.range_ops_loaded = true;
-                    }
-                    let mut code = match bin.op.kind {
-                        TokenKind::Closed => "ClosedRange(",
-                        TokenKind::LeftOpen => "LeftOpenRange(",
-                        TokenKind::RightOpen => "RightOpenRange(",
-                        TokenKind::Open => "OpenRange(",
-                        _ => unreachable!(),
-                    }
-                    .to_string();
-                    code += &self.transpile_expr(*bin.lhs);
-                    code.push(',');
-                    code += &self.transpile_expr(*bin.rhs);
-                    code.push(')');
-                    code
-                }
-                TokenKind::InOp => {
-                    if !self.in_op_loaded {
-                        self.load_in_op();
-                        self.in_op_loaded = true;
-                    }
-                    let mut code = "in_operator(".to_string();
-                    code += &self.transpile_expr(*bin.lhs);
-                    code.push(',');
-                    code += &self.transpile_expr(*bin.rhs);
-                    code.push(')');
-                    code
-                }
-                _ => {
-                    let mut code = "(".to_string();
-                    code += &self.transpile_expr(*bin.lhs);
-                    code.push(' ');
-                    code += &bin.op.content;
-                    code.push(' ');
-                    code += &self.transpile_expr(*bin.rhs);
-                    code += ")";
-                    code
-                }
-            },
             Expr::UnaryOp(unary) => {
                 let mut code = "(".to_string();
                 if unary.op.kind != TokenKind::Mutate {
@@ -344,31 +323,7 @@ impl ScriptGenerator {
                 }
                 other => todo!("transpiling {other}"),
             },
-            Expr::Record(rec) => {
+            Expr::Record(rec) => self.transpile_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 += ")";
-                format!("NamedTuple__('Record', {attrs}){values}")
-            }
             Expr::Tuple(tuple) => match tuple {
                 Tuple::Normal(tup) => {
                     let mut code = "(".to_string();
@@ -394,28 +349,11 @@ impl ScriptGenerator {
                 }
                 other => todo!("transpiling {other}"),
             },
-            Expr::Accessor(acc) => match acc {
+            Expr::Accessor(acc) => self.transpile_acc(acc),
-                Accessor::Ident(ident) => {
-                    match &ident.inspect()[..] {
-                        "Str" | "Bool" | "Nat" | "Array" if !self.builtin_types_loaded => {
-                            self.load_builtin_types();
-                            self.builtin_types_loaded = true;
-                        }
-                        _ => {}
-                    }
-                    Self::transpile_ident(ident)
-                }
-                Accessor::Attr(attr) => {
-                    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::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 {
@@ -449,7 +387,105 @@ impl ScriptGenerator {
                 );
                 String::new()
             }
-            other => todo!("transpile {other:?}"),
+            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;
+                }
+                let mut code = match bin.op.kind {
+                    TokenKind::Closed => "ClosedRange(",
+                    TokenKind::LeftOpen => "LeftOpenRange(",
+                    TokenKind::RightOpen => "RightOpenRange(",
+                    TokenKind::Open => "OpenRange(",
+                    _ => unreachable!(),
+                }
+                .to_string();
+                code += &self.transpile_expr(*bin.lhs);
+                code.push(',');
+                code += &self.transpile_expr(*bin.rhs);
+                code.push(')');
+                code
+            }
+            TokenKind::InOp => {
+                if !self.in_op_loaded {
+                    self.load_in_op();
+                    self.in_op_loaded = true;
+                }
+                let mut code = "in_operator(".to_string();
+                code += &self.transpile_expr(*bin.lhs);
+                code.push(',');
+                code += &self.transpile_expr(*bin.rhs);
+                code.push(')');
+                code
+            }
+            _ => {
+                let mut code = "(".to_string();
+                code += &self.transpile_expr(*bin.lhs);
+                code.push(' ');
+                code += &bin.op.content;
+                code.push(' ');
+                code += &self.transpile_expr(*bin.rhs);
+                code += ")";
+                code
+            }
+        }
+    }
+
+    fn transpile_acc(&mut self, acc: Accessor) -> String {
+        match acc {
+            Accessor::Ident(ident) => {
+                match &ident.inspect()[..] {
+                    "Str" | "Bool" | "Nat" | "Array" if !self.builtin_types_loaded => {
+                        self.load_builtin_types();
+                        self.builtin_types_loaded = true;
+                    }
+                    _ => {}
+                }
+                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)
+                    )
+                }
+            }
         }
     }
 
@@ -462,52 +498,8 @@ impl ScriptGenerator {
                 }
                 code
             }
-            Some("if" | "if!") => {
+            Some("not") => format!("(not ({}))", self.transpile_expr(call.args.remove(0))),
-                let cond = self.transpile_expr(call.args.remove(0));
+            Some("if" | "if!") => self.transpile_if(call),
-                let Expr::Lambda(mut then_block) = call.args.remove(0) else { todo!() };
-                let else_block = call.args.try_remove(0).map(|ex| {
-                    if let Expr::Lambda(blk) = ex {
-                        blk
-                    } else {
-                        todo!()
-                    }
-                });
-                if then_block.body.len() == 1
-                    && else_block
-                        .as_ref()
-                        .map(|blk| blk.body.len() == 1)
-                        .unwrap_or(true)
-                {
-                    let then = self.transpile_expr(then_block.body.remove(0));
-                    if let Some(mut else_block) = else_block {
-                        let els = self.transpile_expr(else_block.body.remove(0));
-                        return format!("{then} if {cond} else {els}");
-                    } else {
-                        return format!("{then} if {cond} else None");
-                    }
-                }
-                let tmp = Str::from(format!("if_tmp_{}__", self.fresh_var_n));
-                self.fresh_var_n += 1;
-                let tmp_func = Str::from(format!("if_tmp_func_{}__", self.fresh_var_n));
-                self.fresh_var_n += 1;
-                let mut code = format!("def {tmp_func}():\n");
-                code += &"    ".repeat(self.level);
-                code += "if ";
-                code += &format!("{cond}:\n");
-                log!(err "{then_block}");
-                code += &self.transpile_block(then_block.body, StoreTmp(tmp.clone()));
-                if let Some(else_block) = else_block {
-                    code += &"    ".repeat(self.level);
-                    code += "else:\n";
-                    code += &self.transpile_block(else_block.body, StoreTmp(tmp.clone()));
-                }
-                code += &"    ".repeat(self.level);
-                code += &format!("return {tmp}\n");
-                self.prelude += &code;
-                // NOTE: In Python, the variable environment of a function is determined at call time
-                // 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();
                 let iter = call.args.remove(0);
@@ -527,60 +519,118 @@ impl ScriptGenerator {
                 code += &self.transpile_block(block.body, Discard);
                 code
             }
-            Some("match" | "match!") => {
+            Some("match" | "match!") => self.transpile_match(call),
-                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));
-                self.fresh_var_n += 1;
-                let mut code = format!("def {tmp_func}():\n");
-                self.level += 1;
-                code += &"    ".repeat(self.level);
-                code += "match ";
-                let cond = call.args.remove(0);
-                code += &format!("{}:\n", self.transpile_expr(cond));
-                while let Some(Expr::Lambda(arm)) = call.args.try_remove(0) {
-                    self.level += 1;
-                    code += &"    ".repeat(self.level);
-                    let target = arm.params.non_defaults.get(0).unwrap();
-                    match &target.pat {
-                        ParamPattern::VarName(param) => {
-                            code += &format!("case {}__:\n", &param.token().content);
-                            code += &self.transpile_block(arm.body, StoreTmp(tmp.clone()));
-                            self.level -= 1;
-                        }
-                        ParamPattern::Discard(_) => {
-                            match target.t_spec.as_ref().map(|t| &t.t_spec) {
-                                Some(TypeSpec::Enum(enum_t)) => {
-                                    let values = ValueObj::vec_from_const_args(enum_t.clone());
-                                    if values.len() == 1 {
-                                        code += &format!("case {}:\n", values[0]);
-                                    } else {
-                                        todo!()
-                                    }
-                                }
-                                Some(_) => todo!(),
-                                None => {
-                                    code += "case _:\n";
-                                }
-                            }
-                            code += &self.transpile_block(arm.body, StoreTmp(tmp.clone()));
-                            self.level -= 1;
-                        }
-                        _ => todo!(),
-                    }
-                }
-                code += &"    ".repeat(self.level);
-                code += &format!("return {tmp}\n");
-                self.prelude += &code;
-                format!("{tmp_func}()")
-            }
             _ => self.transpile_simple_call(call),
         }
     }
 
-    fn transpile_simple_call(&mut self, mut call: Call) -> String {
+    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| {
+            if let Expr::Lambda(blk) = ex {
+                blk
+            } else {
+                todo!()
+            }
+        });
+        if then_block.body.len() == 1
+            && else_block
+                .as_ref()
+                .map(|blk| blk.body.len() == 1)
+                .unwrap_or(true)
+        {
+            let then = self.transpile_expr(then_block.body.remove(0));
+            if let Some(mut else_block) = else_block {
+                let els = self.transpile_expr(else_block.body.remove(0));
+                return format!("{then} if {cond} else {els}");
+            } else {
+                return format!("{then} if {cond} else None");
+            }
+        }
+        let tmp = Str::from(format!("if_tmp_{}__", self.fresh_var_n));
+        self.fresh_var_n += 1;
+        let tmp_func = Str::from(format!("if_tmp_func_{}__", self.fresh_var_n));
+        self.fresh_var_n += 1;
+        let mut code = format!("def {tmp_func}():\n");
+        code += &"    ".repeat(self.level);
+        code += "if ";
+        code += &format!("{cond}:\n");
+        log!(err "{then_block}");
+        code += &self.transpile_block(then_block.body, StoreTmp(tmp.clone()));
+        if let Some(else_block) = else_block {
+            code += &"    ".repeat(self.level);
+            code += "else:\n";
+            code += &self.transpile_block(else_block.body, StoreTmp(tmp.clone()));
+        }
+        code += &"    ".repeat(self.level);
+        code += &format!("return {tmp}\n");
+        self.prelude += &code;
+        // NOTE: In Python, the variable environment of a function is determined at call time
+        // This is a very bad design, but can be used for this code
+        format!("{tmp_func}()")
+    }
+
+    fn transpile_match(&mut self, mut call: Call) -> String {
+        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));
+        self.fresh_var_n += 1;
+        let mut code = format!("def {tmp_func}():\n");
+        self.level += 1;
+        code += &"    ".repeat(self.level);
+        code += "match ";
+        let cond = call.args.remove(0);
+        code += &format!("{}:\n", self.transpile_expr(cond));
+        while let Some(Expr::Lambda(arm)) = call.args.try_remove(0) {
+            self.level += 1;
+            code += &"    ".repeat(self.level);
+            let target = arm.params.non_defaults.get(0).unwrap();
+            match &target.pat {
+                ParamPattern::VarName(param) => {
+                    code += &format!("case {}__:\n", &param.token().content);
+                    code += &self.transpile_block(arm.body, StoreTmp(tmp.clone()));
+                    self.level -= 1;
+                }
+                ParamPattern::Discard(_) => {
+                    match target.t_spec.as_ref().map(|t| &t.t_spec) {
+                        Some(TypeSpec::Enum(enum_t)) => {
+                            let values = ValueObj::vec_from_const_args(enum_t.clone());
+                            if values.len() == 1 {
+                                code += &format!("case {}:\n", values[0]);
+                            } else {
+                                todo!()
+                            }
+                        }
+                        Some(_) => todo!(),
+                        None => {
+                            code += "case _:\n";
+                        }
+                    }
+                    code += &self.transpile_block(arm.body, StoreTmp(tmp.clone()));
+                    self.level -= 1;
+                }
+                _ => todo!(),
+            }
+        }
+        code += &"    ".repeat(self.level);
+        code += &format!("return {tmp}\n");
+        self.prelude += &code;
+        format!("{tmp_func}()")
+    }
+
+    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.push('(');
+        code += &self.transpile_args(call.args, is_py_api, true);
-        while let Some(arg) = call.args.try_remove_pos(0) {
+        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) = 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)
             );
         }
-        code.push(')');
+        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
+    }
 }