diff --git a/src/codegen.hpp b/src/codegen.hpp
index f0c1dba..93b8d90 100644
--- a/src/codegen.hpp
+++ b/src/codegen.hpp
@@ -75,7 +75,7 @@ private:
                 }
                 indent_level--;
             }
-
+            write_indent();
             write("end\n");
         }
         else if (auto *ret = dynamic_cast<const Return *>(&stmt))
@@ -100,6 +100,7 @@ private:
                 generate_stmt(*s);
             }
             indent_level--;
+            write_indent();
             write("end\n");
         }
         else if (auto *while_stmt = dynamic_cast<const While *>(&stmt))
@@ -114,6 +115,7 @@ private:
                 generate_stmt(*s);
             }
             indent_level--;
+            write_indent();
             write("end\n");
         }
         else if (auto *pass_stmt = dynamic_cast<const Pass *>(&stmt))
@@ -149,6 +151,7 @@ private:
                 generate_stmt(*s);
             }
             indent_level--;
+            write_indent();
             write("end\n");
         }
     }