Fix #310

compiler/erg_compiler/codegen.rs

@@ -354,25 +354,15 @@ impl PyCodeGenerator {
     // [..., EXTENDED_ARG 0, EXTENDED_ARG 0, EXTENDED_ARG 1, JUMP_ABSOLUTE 14]
     #[inline]
     fn extend_arg(&mut self, before_instr: usize, bytes: &[u8]) {
+        for byte in bytes.iter().rev().skip(1) {
             self.mut_cur_block_codeobj()
                 .code
-            .insert(before_instr, bytes[2]);
+                .insert(before_instr, *byte);
             self.mut_cur_block_codeobj()
                 .code
                 .insert(before_instr, CommonOpcode::EXTENDED_ARG as u8);
-        self.mut_cur_block_codeobj()
-            .code
-            .insert(before_instr, bytes[1]);
-        self.mut_cur_block_codeobj()
-            .code
-            .insert(before_instr, CommonOpcode::EXTENDED_ARG as u8);
-        self.mut_cur_block_codeobj()
-            .code
-            .insert(before_instr, bytes[0]);
-        self.mut_cur_block_codeobj()
-            .code
-            .insert(before_instr, CommonOpcode::EXTENDED_ARG as u8);
-        self.mut_cur_block().lasti += 6;
+            self.mut_cur_block().lasti += 2;
+        }
     }
 
     fn write_instr<C: Into<u8>>(&mut self, code: C) {
@@ -388,16 +378,40 @@ impl PyCodeGenerator {
                 self.mut_cur_block().lasti += 1;
                 // log!(info "wrote: {}", code);
             }
+            Err(_) => match u16::try_from(code) {
+                Ok(_) => {
+                    let delta =
+                        if CommonOpcode::is_jump_op(*self.cur_block_codeobj().code.last().unwrap())
+                        {
+                            let shift_bytes = 2;
+                            self.jump_delta(code) + shift_bytes
+                        } else {
+                            0
+                        };
+                    let arg = code + delta;
+                    let bytes = u16::try_from(arg).unwrap().to_be_bytes(); // [u8; 2]
+                    let before_instr = self.lasti().saturating_sub(1);
+                    self.mut_cur_block_codeobj().code.push(bytes[1]);
+                    self.mut_cur_block().lasti += 1;
+                    self.extend_arg(before_instr, &bytes);
+                }
                 Err(_) => {
-                let delta = self.jump_delta(code);
+                    let delta =
+                        if CommonOpcode::is_jump_op(*self.cur_block_codeobj().code.last().unwrap())
+                        {
                             let shift_bytes = 6;
-                let arg = code + delta + shift_bytes;
-                let bytes = u32::try_from(arg).unwrap().to_be_bytes();
+                            self.jump_delta(code) + shift_bytes
+                        } else {
+                            0
+                        };
+                    let arg = code + delta;
+                    let bytes = u32::try_from(arg).unwrap().to_be_bytes(); // [u8; 4]
                     let before_instr = self.lasti().saturating_sub(1);
                     self.mut_cur_block_codeobj().code.push(bytes[3]);
                     self.mut_cur_block().lasti += 1;
                     self.extend_arg(before_instr, &bytes);
                 }
+            },
         }
     }
 

tests/should_ok/extended_arg.er

@@ -1,55 +0,0 @@
-a = 1 + 1 + 1
-b = 1 + 1 + 1
-c = 1 + 1 + 1
-d = 1 + 1 + 1
-e = 1 + 1 + 1
-f = 1 + 1 + 1
-g = 1 + 1 + 1
-h = 1 + 1 + 1
-i = 1 + 1 + 1
-j = 1 + 1 + 1
-k = 1 + 1 + 1
-l = 1 + 1 + 1
-m = 1 + 1 + 1
-n = 1 + 1 + 1
-o = 1 + 1 + 1
-p = 1 + 1 + 1
-q = 1 + 1 + 1
-r = 1 + 1 + 1
-s = 1 + 1 + 1
-t = 1 + 1 + 1
-u = 1 + 1 + 1
-v = 1 + 1 + 1
-w = 1 + 1 + 1
-x = 1 + 1 + 1
-y = 1 + 1 + 1
-z = 1 + 1 + 1
-aa = 1 + 1 + 1
-ab = 1 + 1 + 1
-ac = 1 + 1 + 1
-ad = 1 + 1 + 1
-ae = 1 + 1 + 1
-af = 1 + 1 + 1
-ag = 1 + 1 + 1
-ah = 1 + 1 + 1
-ai = 1 + 1 + 1
-aj = 1 + 1 + 1
-ak = 1 + 1 + 1
-al = 1 + 1 + 1
-am = 1 + 1 + 1
-an = 1 + 1 + 1
-ao = 1 + 1 + 1
-ap = 1 + 1 + 1
-aq = 1 + 1 + 1
-ar = 1 + 1 + 1
-as = 1 + 1 + 1
-at = 1 + 1 + 1
-au = 1 + 1 + 1
-av = 1 + 1 + 1
-aw = 1 + 1 + 1
-ax = 1 + 1 + 1
-ay = 1 + 1 + 1
-az = 1 + 1 + 1
-
-assert az == 3
-print! az

tests/should_ok/long.er

@@ -0,0 +1,265 @@
+print! "start"
+x_1 = 1
+x_2 = 1
+x_3 = 1
+x_4 = 1
+x_5 = 1
+x_6 = 1
+x_7 = 1
+x_8 = 1
+x_9 = 1
+x_10 = 1
+x_11 = 1
+x_12 = 1
+x_13 = 1
+x_14 = 1
+x_15 = 1
+x_16 = 1
+x_17 = 1
+x_18 = 1
+x_19 = 1
+x_20 = 1
+x_21 = 1
+x_22 = 1
+x_23 = 1
+x_24 = 1
+x_25 = 1
+x_26 = 1
+x_27 = 1
+x_28 = 1
+x_29 = 1
+x_30 = 1
+x_31 = 1
+x_32 = 1
+x_33 = 1
+x_34 = 1
+x_35 = 1
+x_36 = 1
+x_37 = 1
+x_38 = 1
+x_39 = 1
+x_40 = 1
+x_41 = 1
+x_42 = 1
+x_43 = 1
+x_44 = 1
+x_45 = 1
+x_46 = 1
+x_47 = 1
+x_48 = 1
+x_49 = 1
+x_50 = 1
+x_51 = 1
+x_52 = 1
+x_53 = 1
+x_54 = 1
+x_55 = 1
+x_56 = 1
+x_57 = 1
+x_58 = 1
+x_59 = 1
+x_60 = 1
+x_61 = 1
+x_62 = 1
+x_63 = 1
+x_64 = 1
+x_65 = 1
+x_66 = 1
+x_67 = 1
+x_68 = 1
+x_69 = 1
+x_70 = 1
+x_71 = 1
+x_72 = 1
+x_73 = 1
+x_74 = 1
+x_75 = 1
+x_76 = 1
+x_77 = 1
+x_78 = 1
+x_79 = 1
+x_80 = 1
+x_81 = 1
+x_82 = 1
+x_83 = 1
+x_84 = 1
+x_85 = 1
+x_86 = 1
+x_87 = 1
+x_88 = 1
+x_89 = 1
+x_90 = 1
+x_91 = 1
+x_92 = 1
+x_93 = 1
+x_94 = 1
+x_95 = 1
+x_96 = 1
+x_97 = 1
+x_98 = 1
+x_99 = 1
+x_100 = 1
+x_101 = 1
+x_102 = 1
+x_103 = 1
+x_104 = 1
+x_105 = 1
+x_106 = 1
+x_107 = 1
+x_108 = 1
+x_109 = 1
+x_110 = 1
+x_111 = 1
+x_112 = 1
+x_113 = 1
+x_114 = 1
+x_115 = 1
+x_116 = 1
+x_117 = 1
+x_118 = 1
+x_119 = 1
+x_120 = 1
+x_121 = 1
+x_122 = 1
+x_123 = 1
+x_124 = 1
+x_125 = 1
+x_126 = 1
+x_127 = 1
+x_128 = 1
+x_129 = 1
+x_130 = 1
+x_131 = 1
+x_132 = 1
+x_133 = 1
+x_134 = 1
+x_135 = 1
+x_136 = 1
+x_137 = 1
+x_138 = 1
+x_139 = 1
+x_140 = 1
+x_141 = 1
+x_142 = 1
+x_143 = 1
+x_144 = 1
+x_145 = 1
+x_146 = 1
+x_147 = 1
+x_148 = 1
+x_149 = 1
+x_150 = 1
+x_151 = 1
+x_152 = 1
+x_153 = 1
+x_154 = 1
+x_155 = 1
+x_156 = 1
+x_157 = 1
+x_158 = 1
+x_159 = 1
+x_160 = 1
+x_161 = 1
+x_162 = 1
+x_163 = 1
+x_164 = 1
+x_165 = 1
+x_166 = 1
+x_167 = 1
+x_168 = 1
+x_169 = 1
+x_170 = 1
+x_171 = 1
+x_172 = 1
+x_173 = 1
+x_174 = 1
+x_175 = 1
+x_176 = 1
+x_177 = 1
+x_178 = 1
+x_179 = 1
+x_180 = 1
+x_181 = 1
+x_182 = 1
+x_183 = 1
+x_184 = 1
+x_185 = 1
+x_186 = 1
+x_187 = 1
+x_188 = 1
+x_189 = 1
+x_190 = 1
+x_191 = 1
+x_192 = 1
+x_193 = 1
+x_194 = 1
+x_195 = 1
+x_196 = 1
+x_197 = 1
+x_198 = 1
+x_199 = 1
+x_200 = 1
+x_201 = 1
+x_202 = 1
+x_203 = 1
+x_204 = 1
+x_205 = 1
+x_206 = 1
+x_207 = 1
+x_208 = 1
+x_209 = 1
+x_210 = 1
+x_211 = 1
+x_212 = 1
+x_213 = 1
+x_214 = 1
+x_215 = 1
+x_216 = 1
+x_217 = 1
+x_218 = 1
+x_219 = 1
+x_220 = 1
+x_221 = 1
+x_222 = 1
+x_223 = 1
+x_224 = 1
+x_225 = 1
+x_226 = 1
+x_227 = 1
+x_228 = 1
+x_229 = 1
+x_230 = 1
+x_231 = 1
+x_232 = 1
+x_233 = 1
+x_234 = 1
+x_235 = 1
+x_236 = 1
+x_237 = 1
+x_238 = 1
+x_239 = 1
+x_240 = 1
+x_241 = 1
+x_242 = 1
+x_243 = 1
+x_244 = 1
+x_245 = 1
+x_246 = 1
+x_247 = 1
+x_248 = 1
+x_249 = 1
+x_250 = 1
+x_251 = 1
+x_252 = 1
+x_253 = 1
+x_254 = 1
+x_255 = 1
+x_256 = 1
+
+for! [0, 1], _i =>
+    print!()
+_ = if!, True, do!:
+    print!()
+while! do! False, do!:
+    print!()
+print! "done"

tests/test.rs

@@ -81,6 +81,11 @@ fn exec_interpolation() -> Result<(), ()> {
     expect_success("tests/should_ok/interpolation.er")
 }
 
+#[test]
+fn exec_long() -> Result<(), ()> {
+    expect_success("tests/should_ok/long.er")
+}
+
 #[test]
 fn exec_mut() -> Result<(), ()> {
     expect_success("examples/mut.er")