DeclarationPrinter.cpp

#include "TypeScript/MLIRLogic/MLIRDeclarationPrinter.h"
#include "TypeScript/MLIRLogic/MLIRPrinter.h"

#include "llvm/ADT/StringSet.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/WithColor.h"
#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "mlir"

namespace typescript
{

    void MLIRDeclarationPrinter::newline()
    {
        os << "\n";
    }

    void MLIRDeclarationPrinter::printBeforeDeclaration()
    {
        os << "@dllimport";
        newline();
    }

    void MLIRDeclarationPrinter::print(mlir::Type type)
    {
        MLIRPrinter mp{};
        mp.printBoxedObjectTypes = true;
        mp.quoteNonIdentifierFieldNames = true;
        mp.getClassSpecialization = getClassSpecialization;
        mp.printType<raw_ostream>(os, type);
    }
    
    void MLIRDeclarationPrinter::printNamespaceBegin(NamespaceInfo::TypePtr elementNamespace) {
        if (elementNamespace && elementNamespace->name.size() > 0)
        {
            os << "namespace " << elementNamespace->fullName << " {";
            newline();
        }
    }

    void MLIRDeclarationPrinter::printNamespaceEnd(NamespaceInfo::TypePtr elementNamespace) {
        if (elementNamespace && elementNamespace->name.size() > 0)
        {
            newline();
            os << "}";
            newline();
        }
    }

    void MLIRDeclarationPrinter::printFloatValue(const APFloat &apValue, raw_ostream &os, bool *printedHex)
    {
        // We would like to output the FP constant value in exponential notation,
        // but we cannot do this if doing so will lose precision.  Check here to
        // make sure that we only output it in exponential format if we can parse
        // the value back and get the same value.
        bool isInf = apValue.isInfinity();
        bool isNaN = apValue.isNaN();
        if (!isInf && !isNaN)
        {
            SmallString<128> strValue;
            apValue.toString(strValue, /*FormatPrecision=*/6, /*FormatMaxPadding=*/0, /*TruncateZero=*/false);

            // Check to make sure that the stringized number is not some string like
            // "Inf" or NaN, that atof will accept, but the lexer will not.  Check
            // that the string matches the "[-+]?[0-9]" regex.
            assert(((strValue[0] >= '0' && strValue[0] <= '9') ||
                    ((strValue[0] == '-' || strValue[0] == '+') &&
                     (strValue[1] >= '0' && strValue[1] <= '9'))) &&
                   "[-+]?[0-9] regex does not match!");

            // Parse back the stringized version and check that the value is equal
            // (i.e., there is no precision loss).
            if (APFloat(apValue.getSemantics(), strValue).bitwiseIsEqual(apValue))
            {
                os << strValue;
                return;
            }

            // If it is not, use the default format of APFloat instead of the
            // exponential notation.
            strValue.clear();
            apValue.toString(strValue);

            // Make sure that we can parse the default form as a float.
            if (strValue.str().contains('.'))
            {
                os << strValue;
                return;
            }
        }

        // Print special values in hexadecimal format. The sign bit should be included
        // in the literal.
        if (printedHex)
            *printedHex = true;
        SmallVector<char, 16> str;
        APInt apInt = apValue.bitcastToAPInt();
        apInt.toString(str, /*Radix=*/16, /*Signed=*/false,
                       /*formatAsCLiteral=*/true);
        os << str;
    }

    void MLIRDeclarationPrinter::printAsFieldName(mlir::Attribute attr)
    {
        mlir::TypeSwitch<mlir::Attribute>(attr)
            .Case<mlir::FlatSymbolRefAttr>([&](auto strAttr)
                                           { os << strAttr.getValue(); })
            .Case<mlir::StringAttr>([&](auto strAttr)
                                    { os << strAttr.getValue(); })
            .Case<mlir::IntegerAttr>([&](auto intAttr) {

                os << "[";

                auto intType = intAttr.getType();
                if (intType.isSignlessInteger(1)) {
                    os << (intAttr.getValue().getBoolValue() ? "true" : "false");
                    // Boolean integer attributes always elides the type.
                    return;
                }

                // Only print attributes as unsigned if they are explicitly unsigned or are
                // signless 1-bit values.  Indexes, signed values, and multi-bit signless
                // values print as signed.
                bool isUnsigned =
                    intType.isUnsignedInteger() || intType.isSignlessInteger(1);
                intAttr.getValue().print(os, !isUnsigned); 

                os << "]";
            })
            .Case<mlir::FloatAttr>([&](auto floatAttr) { 
                os << "[";

                bool printedHex = false;
                printFloatValue(floatAttr.getValue(), os, &printedHex); 

                os << "]";
            })
            .Default([&](auto attr) { os << attr; });
    }

    void MLIRDeclarationPrinter::printAsValue(mlir::Attribute attr)
    {
        mlir::TypeSwitch<mlir::Attribute>(attr)
            .Case<mlir::FlatSymbolRefAttr>([&](auto strAttr)
                                           { os << strAttr.getValue(); })
            .Case<mlir::StringAttr>([&](auto strAttr) {
                os << "\"";
                os.write_escaped(strAttr.getValue());
                os << "\""; 
            })
            .Case<mlir::IntegerAttr>([&](auto intAttr) {
                auto intType = intAttr.getType();
                if (intType.isSignlessInteger(1)) {
                    os << (intAttr.getValue().getBoolValue() ? "true" : "false");
                    // Boolean integer attributes always elides the type.
                    return;
                }

                // Only print attributes as unsigned if they are explicitly unsigned or are
                // signless 1-bit values.  Indexes, signed values, and multi-bit signless
                // values print as signed.
                bool isUnsigned =
                    intType.isUnsignedInteger() || intType.isSignlessInteger(1);
                intAttr.getValue().print(os, !isUnsigned); 
            })
            .Case<mlir::FloatAttr>([&](auto floatAttr) { 
                bool printedHex = false;
                printFloatValue(floatAttr.getValue(), os, &printedHex); 
            })
            .Default([&](auto attr) { os << attr; });
    }

    bool MLIRDeclarationPrinter::filterName(StringRef name)
    {
        return name.starts_with(".");
    }

    bool MLIRDeclarationPrinter::filterField(mlir::Attribute attr)
    {
        return mlir::TypeSwitch<mlir::Attribute, bool>(attr)
            .Case<mlir::FlatSymbolRefAttr>([&](auto strAttr)
                                           { return strAttr.getValue().starts_with("."); })
            .Case<mlir::StringAttr>([&](auto strAttr)
                                    { return strAttr.getValue().starts_with("."); })
            .Default([&](auto attr)
                     { return false; });
    }

    void MLIRDeclarationPrinter::printParams(ArrayRef<mlir::Type> params, mlir::Type thisType)
    {
        os << "(";
        auto separator = false;
        for (auto [index, paramType] : enumerate(params))
        {
            if (index == 0 && paramType == thisType)
                continue;

            if (separator)
                os << ", ";
            separator = true;

            auto actualType = paramType;
            auto isThisType = false;
            auto isConditional = false;
            if (auto optType = dyn_cast<mlir_ts::OptionalType>(paramType))
            {
                isConditional = true;
                actualType = optType.getElementType();
            }
            else if (index == 0 && paramType == thisType)
            {
                isThisType = true;
            }

            os << "p" << index << (isConditional ? "?" : "") << " : ";

            if (isThisType)
            {
                os << "this";
            }
            else
            {
                print(actualType);
            }
        }

        os << ")";
    }

    void MLIRDeclarationPrinter::printFunction(StringRef name, ArrayRef<mlir::Type> params, mlir::Type returnType)
    {
        os << "function " << name;

        printParams(params, mlir::Type());

        if (returnType)
        {
            os << " : ";
            print(returnType);
        }
    }

    void MLIRDeclarationPrinter::printMethod(bool isStatic, StringRef name, ArrayRef<mlir::Type> params, mlir::Type returnType, mlir::Type thisType)
    {
        if (isStatic)
        {
            os << "static ";
        }

        os << name;

        printParams(params, thisType);

        if (returnType)
        {
            os << " : ";
            print(returnType);
        }
    }

    void MLIRDeclarationPrinter::printAccessor(bool isStatic, StringRef keyword, StringRef name, mlir_ts::AccessLevel accessLevel,
                                               ArrayRef<mlir::Type> params, mlir::Type returnType, mlir::Type thisType, StringRef dllName)
    {
        if (!dllName.empty())
        {
            os.indent(4);
            os << "@dllname(\"" << dllName << "\")";
            newline();
        }

        os.indent(4);

        if (accessLevel == mlir_ts::AccessLevel::Protected)
        {
            os << "protected ";
        }
        else if (accessLevel == mlir_ts::AccessLevel::Private)
        {
            os << "private ";
        }

        if (isStatic)
        {
            os << "static ";
        }

        os << keyword << " " << name;

        printParams(params, thisType);

        if (returnType)
        {
            os << " : ";
            print(returnType);
        }

        newline();
    }

    void MLIRDeclarationPrinter::printIndexer(mlir::Type indexType, mlir::Type resultType)
    {
        os.indent(4);

        os << "[p0: ";
        print(indexType);
        os << "]: ";
        print(resultType);
        os << ";";

        newline();
    }

    void MLIRDeclarationPrinter::printTypeDeclaration(StringRef name, NamespaceInfo::TypePtr elementNamespace, mlir::Type type)
    {
        printNamespaceBegin(elementNamespace);

        printBeforeDeclaration();

        os << "type " << name << " = ";
        print(type);
        os << ";";
        newline();

        printNamespaceEnd(elementNamespace);
    }

    void MLIRDeclarationPrinter::printEnum(StringRef name, NamespaceInfo::TypePtr elementNamespace, mlir::DictionaryAttr enumValues)
    {
        printNamespaceBegin(elementNamespace);

        printBeforeDeclaration();

        os << "enum " << name;
        newline();
        os << "{";
        newline();

        for (auto enumValue : enumValues)
        {
            os.indent(4);
            printAsFieldName(enumValue.getName());
            if (enumValue.getValue())
            {
                os << " = ";
                printAsValue(enumValue.getValue());
            }

            os << ",";
            newline();
        }

        os << "}";
        newline();

        printNamespaceEnd(elementNamespace);
    }

    void MLIRDeclarationPrinter::printVariableDeclaration(StringRef name, NamespaceInfo::TypePtr elementNamespace, mlir::Type type, bool isConst, StringRef dllName)
    {
        printNamespaceBegin(elementNamespace);

        printBeforeDeclaration();

        if (!dllName.empty())
        {
            os << "@dllname(\"" << dllName << "\")";
            newline();
        }

        // no TS source syntax expresses "this symbol's storage is a single
        // boxed pointer to the real data, not the data inline" - an inferred
        // (untyped) object-literal export is boxed as ObjectType, but its
        // structural-shape declaration text alone is indistinguishable from an
        // explicitly-typed (unboxed, inline) export of the same shape. Emit a
        // sibling @boxed decorator so the importer's isDynamicImport load
        // (MLIRGenVariables.cpp, the branch that actually runs for every
        // ordinary `import '...'` - see mlirGenImportSharedLib's '.' hack)
        // knows to dereference one extra level instead of reading the tuple
        // inline at the resolved symbol address.
        // The variable's own type is then printed as the bare shape - @boxed already adds the
        // extra dereference - while objects held by reference inside it (a generator returned
        // by a method, say) still print as BoxedObject<...>.
        auto printedType = type;
        if (auto objectType = dyn_cast<mlir_ts::ObjectType>(type))
        {
            auto storageType = objectType.getStorageType();
            if (isa<mlir_ts::TupleType>(storageType) || isa<mlir_ts::ConstTupleType>(storageType) ||
                isa<mlir_ts::ObjectStorageType>(storageType))
            {
                os << "@boxed";
                newline();
                printedType = storageType;
            }
        }

        os << (isConst ? "const" : "let") << " " << name << " : ";
        print(printedType);
        os << ";";
        newline();

        printNamespaceEnd(elementNamespace);
    }

    void MLIRDeclarationPrinter::print(StringRef name, NamespaceInfo::TypePtr elementNamespace, mlir_ts::FunctionType funcType, StringRef dllName)
    {
        printNamespaceBegin(elementNamespace);

        printBeforeDeclaration();

        if (!dllName.empty())
        {
            os << "@dllname(\"" << dllName << "\")";
            newline();
        }

        printFunction(
            name,
            funcType.getParams(),
            funcType.getNumResults() > 0 ? funcType.getResult(0) : mlir::Type());
        os << ";";
        newline();

        printNamespaceEnd(elementNamespace);
    }

    void MLIRDeclarationPrinter::print(ClassInfo::TypePtr classType)
    {
        printNamespaceBegin(classType->elementNamespace);

        printBeforeDeclaration();
        // as layout-relevant as `abstract` below: without it the importer numbers the slots of an
        // ordinary class, `.instanceOf` first
        if (classType->isCppVTable)
        {
            os << "@" CPP_VTABLE;
            newline();
        }

        // `abstract` is layout-relevant, not just semantic: mlirGenClassNew skips
        // synthesizing the `.new` method for abstract classes, so dropping the
        // modifier makes the reimporting module insert a `.new` vtable slot the
        // exporting module never had, shifting every subsequent virtual index and
        // corrupting cross-module virtual dispatch (base-module code calling
        // this.method() lands on the wrong slot).
        if (classType->isAbstract)
        {
            os << "abstract ";
        }

        os << "class " << classType->name;

        if (classType->baseClasses.size() > 0)
        {
            os << " extends ";
            auto any = false;
            for (auto baseClass : classType->baseClasses)
            {
                if (any)
                {
                    os << ", ";
                }

                os << baseClass->fullName;
                any = true;
            }
        }

        if (classType->implements.size() > 0)
        {
            os << " implements ";
            auto any = false;
            for (auto implement : classType->implements)
            {
                if (any) 
                {
                    os << ", ";
                }

                os << implement.interface->fullName;
                any = true;
            }
        }

        newline();
        os << "{";
        newline();

        // Static fields
        for (auto staticField : classType->staticFields)
        {
            if (filterField(staticField.id))
                continue;

            os.indent(4);
            
            if (staticField.accessLevel == mlir_ts::AccessLevel::Protected)
            {
                os << "protected ";
            }
            else if (staticField.accessLevel == mlir_ts::AccessLevel::Private)
            {
                os << "private ";
            }

            os << "static ";
            printAsFieldName(staticField.id);
            os << ": ";
            print(staticField.type);
            os << ";";
            newline();
        }

        // Fields
        auto storageType = cast<mlir_ts::ClassStorageType>(classType->classType.getStorageType());
        for (auto field : storageType.getFields())
        {
            if (filterField(field.id))
                continue;

            // a derived class's storage embeds each base class's storage as a synthetic
            // first field whose id is the base's full name (mlirGenClassHeritageClause);
            // that is memory layout, not a source member - printing it would make the
            // importer parse it as a real extra field, shifting every subsequent field's
            // offset and silently corrupting cross-module field access (the `extends`
            // clause printed above already carries the inheritance).
            if (auto strId = dyn_cast<mlir::StringAttr>(field.id))
            {
                if (llvm::any_of(classType->baseClasses,
                                 [&](auto &baseClass) { return strId.getValue() == baseClass->fullName; }))
                {
                    continue;
                }
            }

            os.indent(4);

            if (field.accessLevel == mlir_ts::AccessLevel::Protected)
            {
                os << "protected ";
            }
            else if (field.accessLevel == mlir_ts::AccessLevel::Private)
            {
                os << "private ";
            }

            printAsFieldName(field.id);
            if (field.isConditional)
                os << "?";
            os << ": ";
            print(field.type);
            os << ";";
            newline();
        }

        // index signature ([x: T]: U;) - not stored in storageType's fields,
        // and (unlike accessors) its get/set methods are ordinary named
        // "get"/"set" methods already covered by the methods loop below; only
        // the signature declaration itself is missing from a plain method
        // reprint, and without it a reimporting module's ClassInfo::indexes
        // stays empty so `obj[i]`/`super[i]` can't resolve (see
        // super-index-access-gap-fix memory).
        for (auto indexInfo : classType->indexes)
        {
            if (!indexInfo.indexSignature || indexInfo.indexSignature.getNumResults() == 0)
                continue;

            printIndexer(indexInfo.indexSignature.getInput(0), indexInfo.indexSignature.getResult(0));
        }

        auto printAccessorGet = [&](auto &accessor) {
            printAccessor(
                accessor.isStatic, "get", accessor.name, accessor.getAccessLevel,
                accessor.get.funcType.getParams(),
                accessor.get.funcType.getNumResults() > 0 ? accessor.get.funcType.getResult(0) : mlir::Type(),
                classType->classType,
                accessor.get.dllName);
        };

        auto printAccessorSet = [&](auto &accessor) {
            printAccessor(
                accessor.isStatic, "set", accessor.name, accessor.setAccessLevel,
                accessor.set.funcType.getParams(),
                mlir::Type(),
                classType->classType,
                accessor.set.dllName);
        };

        // accessor halves printed in the methods loop, where their functions sit
        llvm::StringSet<> printedAccessorFuncs;

        // methods (including static)
        for (auto method : classType->methods)
        {
            if (filterName(method.name))
                continue;

            // get/set accessors are ALSO registered here under their mangled
            // "get_x"/"set_x" funcOp name (so same-file/JIT compiles can find them as
            // ordinary methods) - but that mangled form is printed with real `get`/`set`
            // syntax instead. Printed as a plain method, the importer would register it
            // as an ordinary MethodDeclaration instead of a GetAccessor/SetAccessor, so
            // property-style access (`obj.x`, `super.x`) would never populate the
            // reconstructed class's accessors list and fail to resolve.
            //
            // And printed here, at its place among the methods, not after all of them:
            // the importer builds the vtable in the order it reads the members, and this
            // list is the order the exporting module built it in. With the accessors
            // last, `get textContent()` written before `addText()` swapped their slots,
            // and the importer's `addText` call ran the getter.
            auto accessorPrinted = false;
            for (auto &accessor : classType->accessors)
            {
                if (filterName(accessor.name))
                    continue;

                if (accessor.get && accessor.get.name == method.funcName)
                {
                    printAccessorGet(accessor);
                    printedAccessorFuncs.insert(accessor.get.name);
                    accessorPrinted = true;
                    break;
                }

                if (accessor.set && accessor.set.name == method.funcName)
                {
                    printAccessorSet(accessor);
                    printedAccessorFuncs.insert(accessor.set.name);
                    accessorPrinted = true;
                    break;
                }
            }

            if (accessorPrinted || llvm::any_of(classType->accessors, [&](auto &accessor) {
                    return (accessor.get && accessor.get.name == method.funcName) ||
                           (accessor.set && accessor.set.name == method.funcName);
                }))
            {
                continue;
            }

            if (!method.dllName.empty())
            {
                os.indent(4);
                os << "@dllname(\"" << method.dllName << "\")";
                newline();
            }

            if (method.cppNonTrivialReturn)
            {
                os.indent(4);
                os << "@" CPP_NON_TRIVIAL;
                newline();
            }

            os.indent(4);

            if (method.accessLevel == mlir_ts::AccessLevel::Protected)
            {
                os << "protected ";
            }
            else if (method.accessLevel == mlir_ts::AccessLevel::Private)
            {
                os << "private ";
            }

            if (method.isAbstract)
            {
                os << "abstract ";
            }

            // a @cppvtable method returning a record: as declared, the importer lowers it again
            auto funcType = method.declaredFuncType ? method.declaredFuncType : method.funcType;
            printMethod(
                method.isStatic,
                method.name,
                funcType.getParams(),
                funcType.getNumResults() > 0 ? funcType.getResult(0) : mlir::Type(),
                classType->classType);

            newline();
        }

        // accessors whose functions are not among the methods
        for (auto accessor : classType->accessors)
        {
            if (filterName(accessor.name))
                continue;

            if (accessor.get && !printedAccessorFuncs.contains(accessor.get.name))
            {
                printAccessorGet(accessor);
            }

            if (accessor.set && !printedAccessorFuncs.contains(accessor.set.name))
            {
                printAccessorSet(accessor);
            }
        }

        // TODO: we can't declare generic methods, otherwise we need to declare body of methods
        // generic methods
        // for (auto method : classType->staticGenericMethods)
        // {
        //     if (filterName(method.name)) continue;

        //     os.indent(4);

        //     printMethod(
        //         method.isStatic,
        //         method.name,
        //         method.funcType.getParams(),
        //         method.funcType.getNumResults() > 0 ? method.funcType.getResult(0) : mlir::Type());

        //     newline();
        // }

        os << "}";
        newline();

        printNamespaceEnd(classType->elementNamespace);
    }

    void MLIRDeclarationPrinter::printGenericClass(NamespaceInfo::TypePtr elementNamespace, StringRef rawDeclarationText)
    {
        // unlike print(ClassInfo::TypePtr) above, a generic class has no compiled body for
        // any given instantiation to @dllimport against - each importing module
        // instantiates it locally instead, so the caller passes the ORIGINAL source text
        // (type parameters and method bodies intact) verbatim; this just supplies the
        // matching namespace wrapper, deliberately WITHOUT printBeforeDeclaration()'s
        // @dllimport marker.
        printNamespaceBegin(elementNamespace);
        os << rawDeclarationText;
        newline();
        printNamespaceEnd(elementNamespace);
    }

    void MLIRDeclarationPrinter::print(InterfaceInfo::TypePtr interfaceType)
    {
        printNamespaceBegin(interfaceType->elementNamespace);

        printBeforeDeclaration();
        os << "interface " << interfaceType->name;

        if (interfaceType->extends.size() > 0)
        {
            os << " extends ";
            auto any = false;
            for (auto interfaceInfo : interfaceType->extends)
            {
                if (any) 
                {
                    os << ", ";
                }

                os << interfaceInfo.second->fullName;
                any = true;
            }
        }

        newline();
        os << "{";
        newline();

        // Fields
        for (auto field : interfaceType->fields)
        {
            if (filterField(field.id))
                continue;

            os.indent(4);
            printAsFieldName(field.id);
            if (field.isConditional)
                os << "?";
            os << ": ";
            print(field.type);
            os << ";";
            newline();
        }

        // index signature ([x: T]: U;) - mirrors the class printer's fix
        // (see cross-module-class-indexer-shared-gap-fix memory): without
        // this, a reimporting module's InterfaceInfo::indexes stays empty
        // and `obj[i]` through the interface fails with "indexer is not
        // declared" / "Interface member '.index' can't be found".
        //
        // Unlike a class's index signature (a plain (arg)->result FunctionType),
        // an interface's is built via getInterfaceMethodNameAndType with
        // funcGenContext.thisType set, which prepends an opaque `this` input -
        // so the real index-argument type is input(1), not input(0) (matching
        // getIndexSignatureArgumentAndResultTypes's "first parameter is Opaque"
        // branch in MLIRTypeHelper.h).
        for (auto indexInfo : interfaceType->indexes)
        {
            if (!indexInfo.indexSignature || indexInfo.indexSignature.getNumResults() == 0)
                continue;

            auto argIndex = indexInfo.indexSignature.getNumInputs() > 1 ? 1 : 0;
            printIndexer(indexInfo.indexSignature.getInput(argIndex), indexInfo.indexSignature.getResult(0));
        }

        // methods (including static)
        auto opaqueType = mlir_ts::OpaqueType::get(interfaceType->interfaceType.getContext());
        for (auto method : interfaceType->methods)
        {
            if (filterName(method.name))
                continue;

            os.indent(4);

            printMethod(
                false,
                method.name,
                method.funcType.getParams(),
                method.funcType.getNumResults() > 0 ? method.funcType.getResult(0) : mlir::Type(),
                /*interfaceType->interfaceType*/opaqueType);

            newline();
        }

        // TODO: we can't declare generic methods, otherwise we need to declare body of methods
        // generic methods
        // for (auto method : classType->staticGenericMethods)
        // {
        //     if (filterName(method.name)) continue;

        //     os.indent(4);

        //     printMethod(
        //         method.isStatic,
        //         method.name,
        //         method.funcType.getParams(),
        //         method.funcType.getNumResults() > 0 ? method.funcType.getResult(0) : mlir::Type());

        //     newline();
        // }

        os << "}";
        newline();

        printNamespaceEnd(interfaceType->elementNamespace);
    }    

} // namespace typescript