// Before
func f(x : T) : U { return <expr>; };

// After
func f(x : T) : U { <expr> };

#!/usr/bin/env python3
"""
Remove terminal `return` from Motoko functions that have exactly one return
statement in their direct body (excluding nested functions).

Usage:
  1. Set src_dirs to match your project layout.
  2. Run: python3 remove_returns.py
  3. Run tests: npx mops test
  4. Delete this script after confirming all tests pass.

Safety:
  - Tracks function boundaries via character-level brace-depth parsing.
  - Skips string literals ("..." with \" escapes) and comments (// and /* */).
  - Detects nested `func` declarations and skips their bodies entirely.
  - Only removes the `return` keyword (+ trailing space) from functions with
    exactly 1 return at any depth within the direct body.
"""

import re
import glob
import os

Handles:
  - String literals (skips content inside "...")
  - Line comments (// ...)
  - Block comments (/* ... */)
  - Nested braces
"""
i = 0
n = len(text)
while i < n:
    # Skip string literals
    if text[i] == '"':
        i += 1
        while i < n and text[i] != '"':
            if text[i] == '\\':
                i += 1  # skip escaped char
            i += 1
        i += 1  # skip closing "
        continue

    # Skip line comments
    if text[i] == '/' and i + 1 < n and text[i + 1] == '/':
        i += 2
        while i < n and text[i] != '\n':
            i += 1
        continue

    # Skip block comments
    if text[i] == '/' and i + 1 < n and text[i + 1] == '*':
        i += 2
        while i < n and not (text[i] == '*' and i + 1 < n and text[i + 1] == '/'):
            i += 1
        i += 2  # skip */
        continue

    # Look for 'func' keyword at a word boundary
    if text[i:i+4] == 'func' and (i == 0 or not text[i-1].isalnum() and text[i-1] != '_'):
        after = text[i+4:i+5] if i + 4 < n else ''
        if after == '' or not (after.isalnum() or after == '_'):
            # Found a func keyword. Scan forward to find the opening '{'.
            j = i + 4
            while j < n:
                if text[j] == '"':
                    j += 1
                    while j < n and text[j] != '"':
                        if text[j] == '\\':
                            j += 1
                        j += 1
                    j += 1
                    continue
                if text[j] == '{':
                    # Found the opening brace of the function body.
                    brace_start = j
                    depth = 1
                    j += 1
                    while j < n and depth > 0:
                        if text[j] == '"':
                            j += 1
                            while j < n and text[j] != '"':
                                if text[j] == '\\':
                                    j += 1
                                j += 1
                            j += 1
                            continue
                        if text[j] == '/' and j + 1 < n and text[j + 1] == '/':
                            j += 2
                            while j < n and text[j] != '\n':
                                j += 1
                            continue
                        if text[j] == '/' and j + 1 < n and text[j + 1] == '*':
                            j += 2
                            while j < n and not (text[j] == '*' and j + 1 < n and text[j + 1] == '/'):
                                j += 1
                            j += 2
                            continue
                        if text[j] == '{':
                            depth += 1
                        elif text[j] == '}':
                            depth -= 1
                        j += 1
                    brace_end = j - 1  # index of closing '}'
                    yield (brace_start, brace_end)
                    i = j
                    break
                if text[j] == '=' or text[j] == ';':
                    # func ... = expr; (no body) or forward decl
                    i = j + 1
                    break
                j += 1
            else:
                i = j
            continue
    i += 1

# First, find all nested func bodies within this body so we can skip them.
nested_ranges = []
for ns, ne in find_func_bodies(body):
    # Adjust to absolute positions
    nested_ranges.append((ns + offset, ne + offset))

def is_inside_nested(pos):
    for ns, ne in nested_ranges:
        if ns <= pos <= ne:
            return True
    return False

# Now scan for `return` keywords in the body, skipping nested funcs.
returns = []
i = 0
while i < len(body):
    # Skip strings
    if body[i] == '"':
        i += 1
        while i < len(body) and body[i] != '"':
            if body[i] == '\\':
                i += 1
            i += 1
        i += 1
        continue

    # Skip line comments
    if body[i] == '/' and i + 1 < len(body) and body[i + 1] == '/':
        i += 2
        while i < len(body) and body[i] != '\n':
            i += 1
        continue

    # Skip block comments
    if body[i] == '/' and i + 1 < len(body) and body[i + 1] == '*':
        i += 2
        while i < len(body) and not (body[i] == '*' and i + 1 < len(body) and body[i + 1] == '/'):
            i += 1
        i += 2
        continue

    # Check for 'return' keyword
    if body[i:i+6] == 'return' and (i == 0 or not body[i-1].isalnum() and body[i-1] != '_'):
        after = body[i+6:i+7] if i + 6 < len(body) else ''
        if after == '' or not (after.isalnum() or after == '_'):
            abs_pos = i + offset
            if not is_inside_nested(abs_pos):
                returns.append((abs_pos, abs_pos + 6))
            i += 6
            continue
    i += 1

return returns

original = text
removals = 0

# Collect all function bodies
func_bodies = list(find_func_bodies(text))

# For each function, check if it has exactly 1 return in its direct body
# Process in reverse order to preserve indices when editing
edits = []  # list of (start, end) of "return " to remove

for body_start, body_end in func_bodies:
    returns = count_returns_in_direct_body(text, body_start, body_end)
    if len(returns) == 1:
        ret_start, ret_end = returns[0]
        # Verify this return is the terminal statement of the direct body:
        # scan past the return's expression (tracking brackets/strings/comments)
        # to its terminating ';' (or body_end), then ensure only whitespace,
        # comments, and optional semicolons remain before body_end.
        n = len(text)
        j = ret_end
        depth = 0
        stmt_end = body_end  # position after the return statement
        while j < body_end:
            c = text[j]
            if c == '"':
                j += 1
                while j < body_end and text[j] != '"':
                    if text[j] == '\\':
                        j += 1
                    j += 1
                j += 1
                continue
            if c == '/' and j + 1 < body_end and text[j + 1] == '/':
                j += 2
                while j < body_end and text[j] != '\n':
                    j += 1
                continue
            if c == '/' and j + 1 < body_end and text[j + 1] == '*':
                j += 2
                while j < body_end and not (text[j] == '*' and j + 1 < body_end and text[j + 1] == '/'):
                    j += 1
                j += 2
                continue
            if c in '({[':
                depth += 1
            elif c in ')}]':
                depth -= 1
            elif c == ';' and depth == 0:
                stmt_end = j + 1
                break
            elif c == '\n' and depth == 0:
                # Statement terminated by newline (no semicolon)
                stmt_end = j
                break
            j += 1
        else:
            stmt_end = body_end

        # Now skip whitespace, comments, and stray semicolons after the return
        k = stmt_end
        terminal = True
        while k < body_end:
            c = text[k]
            if c.isspace() or c == ';':
                k += 1
                continue
            if c == '/' and k + 1 < body_end and text[k + 1] == '/':
                k += 2
                while k < body_end and text[k] != '\n':
                    k += 1
                continue
            if c == '/' and k + 1 < body_end and text[k + 1] == '*':
                k += 2
                while k < body_end and not (text[k] == '*' and k + 1 < body_end and text[k + 1] == '/'):
                    k += 1
                k += 2
                continue
            # Found other code after the return — not terminal.
            terminal = False
            break

        if not terminal:
            continue

        # Remove "return " (keyword + trailing space)
        if ret_end < len(text) and text[ret_end] == ' ':
            edits.append((ret_start, ret_end + 1))
        else:
            edits.append((ret_start, ret_end))

# Apply edits in reverse order to preserve positions
edits.sort(key=lambda x: x[0], reverse=True)
for start, end in edits:
    text = text[:start] + text[end:]
    removals += 1

if text != original:
    with open(filepath, 'w') as f:
        f.write(text)

return removals

---

// Before (mixed)
import Runtime "mo:core/Runtime";
import { type Result } "mo:core/Types";
import SHA256 "mo:sha2/SHA256";
import Map "mo:core/Map";
import BitVec "mo:bitvec/BitVec";
import Utils "../lib/Utils";
import Logger "./Logger";

// After (aggregated and sorted per section)
//// Core
import Map "mo:core/Map";
import Runtime "mo:core/Runtime";
import { type Result } "mo:core/Types";

//// Third‑party (mops)
import BitVec "mo:bitvec/BitVec";
import SHA256 "mo:sha2/SHA256";

//// Local
import Logger "Logger";
import Utils "../lib/Utils";

// Before
let b : Blob = Text.encodeUtf8("hello");

// After
let b : Blob = "hello";

// Good: direct assignment
let blobs = [
  "strategy",
  Text.encodeUtf8(Nat.toText(slot)),
];

// Avoid: redundant encoding for constant string
let blobs = [
  Text.encodeUtf8("strategy"),
  Text.encodeUtf8(Nat.toText(slot)),
];

rg -n --glob '!**/.mops/**' --glob '**/*.mo' "func [A-Za-z_][A-Za-z0-9_]*\(.*\) *:.*\{ *return .*; *};"

undefined

4) Aggregate + sort imports into sections (editor macro)
- Select all `import` lines at the top of the file → group into three sections (Core, Third‑party mo:*, Local) → sort each group alphabetically by path → insert blank lines between sections → keep import styles as‑is.

---

// Before
Array.fromVarArray(buf)
Array.fromVarArray<Nat8>(buf)

// After
buf.toArray()