#!/usr/bin/env python3
"""Give every `defn` an explicit return type, and rewrite `Unit` as `()`.

`(defn f [x i32] body)` becomes `(defn f [x i32] () body)`, and a return type
already written as `Unit` -- or a `Unit` anywhere else a type is spelled, as in
`(Fn [i32] Unit)` -- becomes `()`.

Two rules, one pass, because both are the same change: the slot after the
parameters is now unconditionally a type, so a function that returns nothing
has to say so, and the thing it says is `()`.

Deciding whether a `defn` already has a return type is the whole difficulty,
and this script does it the way `lib/parse.ml` did before the slot became
mandatory: a form in that position is the return type when it is a *type form*
and it is not the entire body. `is_type_form` below is a transcription of the
one in parse.ml, deliberately faithful rather than improved -- being identical
to the parser it is replacing is what makes the sweep meaning-preserving. The
type names it needs come from the file's own declarations, from the prelude's
(read out of lib/prelude.ml), and from the builtin list.

Re-runnable: a `defn` whose slot is already filled is left alone, and `()` is
itself a type form, so converting a converted file is a no-op. Parallel
branches that wrote Flan in the old spelling want this pass at merge.

    tools/unit-return.py <file-or-dir>...            # rewrite .flan in place
    tools/unit-return.py --check <file-or-dir>...    # report, change nothing
    tools/unit-return.py --in-strings <file.ml>...   # Flan inside "..." literals
    tools/unit-return.py --raw-ml lib/prelude.ml     # Flan in a {flan|...|flan} block
    tools/unit-return.py --in-html web/index.html    # Flan in <pre><code> blocks

A directory is walked for `.flan` files only. Anything else is named on the
command line with the mode that says how its Flan is embedded, because a blind
scan of an OCaml or HTML file would read its punctuation as Flan.

`-v` logs every `defn` seen and what was decided about it, which is the only
practical way to review a sweep this size.

**Read the diff of every non-`.flan` file.** A snippet split across OCaml
string concatenation -- `decls ^ "(defn f [s [u8]] Cursor ...)"` -- is scanned
one literal at a time, and the names the other literal declared would be
invisible. The embedded modes work around it by pooling every fragment's type
declarations across the whole file, and counting a pooled name only as a bare
symbol, exactly as the prelude's types count: as a list head it would eat
`(Some 1)` and `(Rune {.code 65})` as return types, which is the misparse this
change exists to remove. That is sound because no user type takes arguments --
only the builtin constructors do, and they are known already -- but it is a
pool and not the real scope, so read the diff.

And one thing it cannot know at all: a file that spells the *refused* forms on
purpose, to test that they are refused. `test/test_flan.ml` holds six such
sites -- `(defn f [] (g))`, `(defn f [] Unit (g))`, `(defn f [] Nope (bar))`,
`(defn f [] f65 0.0)` -- and this script encodes the old rule, so it wants to
convert every one of them and must not. A clean run over this repo reports
exactly those six; anything else is a real conversion.
"""

import sys, os, re

DELIM = set('()[]{}";`~ \t\n\r,')
OPENERS = {'(': ')', '[': ']', '{': '}'}
CLOSERS = {')', ']', '}'}

# lib/parse.ml, [primitives] and [builtin_types]. Present in [types] under
# their plain names, so they count as type forms in every position.
BUILTINS = {
    "i8", "i16", "i32", "i64", "u8", "u16", "u32", "u64",
    "f32", "f64", "bool", "string", "Unit", "Never",
    "Ptr", "Option", "Result", "Vec", "Map", "Handle", "Fn",
}

NUMERIC = re.compile(r'^[-+]?[0-9]')


class Atom:
    def __init__(self, start, end, text):
        self.start, self.end, self.text = start, end, text
        self.tok = start        # never moves; `start` may slide onto a quote sigil
        self.items = []

    @property
    def open_char(self):
        return None

    def sym(self):
        """The symbol name, or None if this atom is not one.

        A string, a keyword, a character literal and a number are all atoms and
        none of them is a symbol -- parse.ml's [is_type_form] answers false for
        every one of them, through its catch-all arm.
        """
        t = self.text
        if not t or t[0] in '":\\' or NUMERIC.match(t):
            return None
        return t


class Seq:
    def __init__(self, open_char, start):
        self.open_char, self.start = open_char, start
        self.end = start
        self.items = []
        self.closed = False

    def sym(self):
        return None


def lex_forms(src, i, end, stop=None):
    """Read forms from src[i:end]. Returns (items, next_index, closed).

    `closed` is False when the text ran out before the enclosing delimiter did,
    which is how a fragment that holds only part of a form is recognised --
    `"(defn step [] i64\\n"`, one line of a snippet built by concatenation.

    The lexing rules mirror lib/reader.ml. A quote/quasiquote/unquote prefix is
    folded into the form it applies to, so a quoted value stays one element.
    """
    items = []
    n = end
    pending_prefix = None

    def push(node):
        nonlocal pending_prefix
        if pending_prefix is not None:
            node.start = pending_prefix
            pending_prefix = None
        items.append(node)

    while i < n:
        c = src[i]
        if c in ' \t\n\r,':
            i += 1
        elif c == ';':                                   # line comment
            while i < n and src[i] != '\n':
                i += 1
        elif c == '"':                                   # string literal
            j = i + 1
            while j < n and src[j] != '"':
                j += 2 if src[j] == '\\' else 1
            j = min(j + 1, n)
            push(Atom(i, j, src[i:j]))
            i = j
        elif c == '\\':                                  # character literal
            j = i + 1
            if j < n:
                j += 1
                while j < n and src[j] not in DELIM:
                    j += 1
            push(Atom(i, j, src[i:j]))
            i = j
        elif c in "'`~":                                 # quote sugar
            if pending_prefix is None:
                pending_prefix = i
            i += 2 if (c == '~' and i + 1 < n and src[i + 1] == '@') else 1
        elif c in OPENERS:
            node = Seq(c, i)
            node.items, i, node.closed = lex_forms(src, i + 1, n, OPENERS[c])
            node.end = i
            push(node)
        elif c in CLOSERS:
            return items, i + 1, True
        else:                                            # symbol or keyword
            j = i
            while j < n and src[j] not in DELIM:
                j += 1
            if j == i:
                j = i + 1
            push(Atom(i, j, src[i:j]))
            i = j
    return items, n, False


def head(node):
    """The head symbol of a `(...)` form, or None."""
    if getattr(node, 'open_char', None) != '(' or not node.items:
        return None
    return node.items[0].sym()


# ── the type sets, as lib/parse.ml collects them ──────────────────────

class Types:
    """What `is_type_form` consults.

    Three sets, kept apart exactly as parse.ml keeps them, because which
    positions a name counts in depends on where it came from:

    `names` -- builtins and the file's own defstruct/defunion/defalias. Count
    as a bare symbol *and* as a list head, since `(Option f64)` is a type.

    `enums`, `prelude` -- enum names and the prelude's types. Count only as a
    bare symbol. As a list head they would eat `(Key 1)` and `(Rune {.code 65})`
    -- a conversion and a constructor -- as return types, which is the silent
    misparse this whole change exists to remove.

    `aliases` -- import aliases, for `rl/Vector2`.
    """
    def __init__(self, prelude_names, prelude_enums):
        self.names = set(BUILTINS)
        self.enums = set(prelude_enums)
        self.prelude = set(prelude_names)
        self.aliases = set()

    def scan(self, forms):
        """Add what a program's top-level declarations introduce."""
        for f in forms:
            h = head(f)
            if h in ('defstruct', 'defunion', 'defalias') and len(f.items) == 3:
                n = f.items[1].sym()
                if n:
                    self.names.add(n)
            elif h == 'defenum' and len(f.items) == 3:
                n = f.items[1].sym()
                if n:
                    self.enums.add(n)
            elif h == 'import' and len(f.items) == 3:
                a = f.items[1].sym()
                if a:
                    self.aliases.add(a)

    def copy(self):
        t = Types(self.prelude, self.enums)
        t.names = set(self.names)
        t.aliases = set(self.aliases)
        return t

    def qualified(self, s):
        i = s.find('/')
        if i < 0:
            return False
        alias, name = s[:i], s[i + 1:]
        return alias in self.aliases and name[:1].isupper()

    def is_type_form(self, f):
        oc = getattr(f, 'open_char', None)
        if oc == '(' and not f.items:
            return True                      # () is unit
        if oc in ('[', '{'):
            return True                      # [T], [n T] and {K V} are only types
        if oc == '(':
            h = head(f)
            return bool(h) and (h in self.names or self.qualified(h))
        s = f.sym()
        if s is None:
            return False
        return s in self.names or s in self.enums or s in self.prelude \
            or self.qualified(s)


def prelude_type_names(root):
    """The prelude's type names, from its {flan|...|flan} block.

    parse.ml's [prelude_types] does the same walk over the same text; reading
    the file keeps the two from drifting apart by hand.
    """
    path = os.path.join(root, 'lib', 'prelude.ml')
    try:
        with open(path, encoding='utf-8') as fh:
            src = fh.read()
    except OSError:
        return set(), set()
    body = flan_block(src)
    if body is None:
        return set(), set()
    forms, _, _ = lex_forms(body, 0, len(body))
    names, enums = set(), set()
    for f in forms:
        h = head(f)
        if len(f.items) != 3:
            continue
        n = f.items[1].sym()
        if not n:
            continue
        if h in ('defstruct', 'defunion', 'defalias'):
            names.add(n)
        elif h == 'defenum':
            enums.add(n)
    return names, enums


def flan_block(src):
    i = src.find('{flan|')
    j = src.rfind('|flan}')
    if i < 0 or j < i:
        return None
    return src[i + len('{flan|'):j]


# ── the rewrite ───────────────────────────────────────────────────────

def plan(src, types, label, log, base_line=1):
    """Edits for one program's worth of Flan. Returns [(start, end, text)]."""
    forms, _, _ = lex_forms(src, 0, len(src))
    types = types.copy()
    types.scan(forms)
    edits = []

    def line_of(off):
        return base_line + src.count('\n', 0, off)

    def walk(node, nested):
        if getattr(node, 'open_char', None) is None:
            # Only inside a form. Every type position is -- (Fn [i32] Unit),
            # [Unit], the slot after a defn's parameters -- and a bare top-level
            # `Unit` is not Flan at all. The guard is what keeps this pass off
            # an OCaml literal that happens to spell the word, as the checker's
            # own pattern `Tname "Unit"` does.
            if nested and node.sym() == 'Unit':
                edits.append((node.tok, node.end, '()'))
            return
        if head(node) == 'defn':
            decide(node)
        for it in node.items:
            walk(it, True)

    def decide(node):
        items = node.items
        # (defn name [params] ...). Anything else -- a metadata sigil, a
        # malformed form -- is left alone and reported, because guessing at a
        # shape the parser does not accept is how a sweep corrupts a file.
        if not node.closed:
            log.append("%s:%d: skipped, the form is cut off here -- a fragment "
                       "of a snippet built by concatenation"
                       % (label, line_of(node.start)))
            return
        if len(items) < 3 or getattr(items[2], 'open_char', None) != '[':
            log.append("%s:%d: skipped, not (defn name [params] ...)"
                       % (label, line_of(node.start)))
            return
        name = items[1].sym() or '?'
        rest = items[3:]
        # parse.ml's guard: a single remaining form is the body, not the
        # return type -- [(defn f [] i32)] was a function returning Unit whose
        # body is the name [i32]. The exception is a lone [()], which is what
        # this script itself writes for a function with no body, and which was
        # never a legal body form. Without it a second run would fill the slot
        # again, and a re-runnable sweep is the point.
        lone_unit = len(rest) == 1 and getattr(rest[0], 'open_char', None) == '(' \
            and not rest[0].items
        if rest and (len(rest) >= 2 or lone_unit) and types.is_type_form(rest[0]):
            log.append("%s:%d: %s kept %s"
                       % (label, line_of(node.start), name,
                          src[rest[0].start:rest[0].end].replace('\n', ' ')))
            return
        edits.append((items[2].end, items[2].end, ' ()'))
        log.append("%s:%d: %s filled ()" % (label, line_of(node.start), name))

    for f in forms:
        walk(f, False)
    return edits


def apply(src, edits):
    if not edits:
        return src, 0
    out = []
    last = 0
    for start, end, text in sorted(edits):
        out.append(src[last:start])
        out.append(text)
        last = end
    out.append(src[last:])
    return ''.join(out), len(edits)


def mask_ocaml_escapes(body):
    """Blank out OCaml escapes so the Flan lexer cannot trip on them.

    Length-preserving, so offsets into the masked text index the original. A
    `\\"` must not end a Flan string, and the `\\n\\` line continuations the
    tests wrap their snippets with must not read as Flan character literals.

    An escaped newline becomes a real one rather than a blank: a `;` comment
    runs to end of line, so flattening `\\n` to spaces would let one comment
    swallow the rest of the snippet.
    """
    chars = list(body)
    i = 0
    while i < len(chars) - 1:
        if chars[i] == '\\':
            c = chars[i + 1]
            chars[i] = ' '
            chars[i + 1] = c if c in '\n\t' else ('\n' if c == 'n' else ' ')
            i += 2
        else:
            i += 1
    return ''.join(chars)


def pool(types, fragments):
    """Fold every fragment's declarations into the bare-symbol-only sets.

    One file's Flan is written in pieces -- concatenated OCaml literals, one
    <pre> per section -- and a piece does not see the piece that declared its
    types. Pooling gives it back. See the module docstring for why pooled names
    count only in bare-symbol position.
    """
    for body in fragments:
        forms, _, _ = lex_forms(body, 0, len(body))
        seen = types.copy()
        seen.names = set()
        seen.enums = set()
        seen.scan(forms)
        types.prelude |= seen.names
        types.enums |= seen.enums
        types.aliases |= seen.aliases
    return types


def convert_flan(src, types, label, log):
    return apply(src, plan(src, types, label, log))


def ocaml_literals(src):
    """(offset, masked body) for every OCaml string literal in `src`."""
    out = []
    i, n = 0, len(src)
    while i < n:
        c = src[i]
        if c == '"':
            j = i + 1
            while j < n and src[j] != '"':
                j += 2 if src[j] == '\\' else 1
            out.append((i + 1, mask_ocaml_escapes(src[i + 1:j])))
            i = j + 1
        elif c == '(' and i + 1 < n and src[i + 1] == '*':   # OCaml comment
            i += 2
        else:
            i += 1
    return out


def convert_in_strings(src, types, label, log):
    """Flan inside ordinary OCaml `"..."` literals, as the tests write it."""
    lits = ocaml_literals(src)
    pool(types, [body for _, body in lits])
    edits = []
    for off, body in lits:
        base = src.count('\n', 0, off) + 1
        for s, e, t in plan(body, types, label, log, base):
            edits.append((off + s, off + e, t))
    return apply(src, edits)


def convert_raw_ml(src, types, label, log):
    """A whole {flan|...|flan} block, as lib/prelude.ml writes it."""
    i = src.find('{flan|')
    j = src.rfind('|flan}')
    if i < 0 or j < i:
        return src, 0
    off = i + len('{flan|')
    body = src[off:j]
    base = src.count('\n', 0, off) + 1
    edits = [(off + s, off + e, t) for s, e, t in plan(body, types, label, log, base)]
    return apply(src, edits)


CODE = re.compile(r'<pre><code>(.*?)</code></pre>', re.S)


def convert_in_html(src, types, label, log):
    """Flan in <pre><code> blocks. HTML entities are left escaped: `&lt;` lexes
    as an ordinary atom and nothing this pass writes needs escaping."""
    edits = []
    pool(types, [m.group(1) for m in CODE.finditer(src)])
    for m in CODE.finditer(src):
        off = m.start(1)
        base = src.count('\n', 0, off) + 1
        for s, e, t in plan(m.group(1), types, label, log, base):
            edits.append((off + s, off + e, t))
    return apply(src, edits)


MODES = {
    '--in-strings': convert_in_strings,
    '--raw-ml': convert_raw_ml,
    '--in-html': convert_in_html,
}


def walk_paths(paths):
    for p in paths:
        if os.path.isdir(p):
            for root, dirs, files in os.walk(p):
                # vendor/ is not third-party: edn, raylib and agent are this
                # repo's own packages, written in Flan, and they convert too.
                dirs[:] = [d for d in dirs
                           if d not in ('_build', '.git', 'node_modules')]
                for f in sorted(files):
                    if f.endswith('.flan'):
                        yield os.path.join(root, f)
        else:
            yield p


def main(argv):
    check = '--check' in argv
    verbose = '-v' in argv or '--verbose' in argv
    modes = [m for m in MODES if m in argv]
    if len(modes) > 1:
        sys.stderr.write("pick one of %s\n" % ', '.join(MODES))
        return 2
    mode = modes[0] if modes else None
    paths = [a for a in argv[1:] if not a.startswith('-')] or ['.']

    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    pnames, penums = prelude_type_names(root)
    if not pnames:
        sys.stderr.write("warning: no prelude types found under %s -- a "
                         "prelude type in return position may be mis-read\n" % root)

    total_files = total_edits = 0
    log = []
    for path in walk_paths(paths):
        if not path.endswith('.flan') and mode is None:
            sys.stderr.write(
                "%s: say how the Flan is embedded -- --in-strings for OCaml "
                "string literals, --raw-ml for a {flan|...|flan} block, "
                "--in-html for <pre><code> blocks. A blind scan would read the "
                "host language's punctuation as Flan.\n" % path)
            return 2
        with open(path, encoding='utf-8') as fh:
            src = fh.read()
        types = Types(pnames, penums)
        fn = MODES[mode] if mode else convert_flan
        new, n = fn(src, types, path, log)
        if n:
            total_files += 1
            total_edits += n
            print("%s: %d" % (path, n))
            if not check:
                with open(path, 'w', encoding='utf-8') as fh:
                    fh.write(new)
    if verbose:
        for line in log:
            print("  " + line)
    verb = "would make" if check else "made"
    print("%s %d edits across %d files" % (verb, total_edits, total_files))
    return 1 if (check and total_edits) else 0


if __name__ == '__main__':
    sys.exit(main(sys.argv))