WoofWare.Myriad/WoofWare.Myriad.Plugins/CataGenerator.fs

namespace WoofWare.Myriad.Plugins

open Fantomas.FCS.Syntax
open Fantomas.FCS.SyntaxTrivia
open Fantomas.FCS.Xml
open WoofWare.Whippet.Fantomas

[<RequireQualifiedAccess>]
module internal CataGenerator =
    open Fantomas.FCS.Text.Range

    /// The user-provided DU contains cases, each of which contains fields.
    /// We have a hard-coded set of things we know how to deal with as field contents.
    type FieldDescription =
        /// type Thing = | Case of Thing list * whatever
        | ListSelf of SynType
        /// type Thing = | Case of Thing * whatever
        | Self of SynType
        /// type Thing = | Case of int * whatever
        | NonRecursive of SynType

    /// Within a union case, there are several fields. This is a field.
    /// (The name is "CataUnionField" merely to distinguish it from the more general
    /// `UnionField` notion we already have in this library; it's got more information
    /// in it that is unique to this source generator.)
    type CataUnionBasicField =
        {
            /// The name of this field as the user originally wrote, if available.
            /// For example, `| Foo of blah : int` would give `Some "blah"`.
            FieldName : Ident option
            /// The name we will use when accessing this field.
            /// This is FieldName if available, or otherwise an autogenerated name.
            ArgName : Ident
            /// The relationship this field has with the parent type (or the
            /// recursive knot of parent types)
            Description : FieldDescription
            /// Any generic parameters this field consumes.
            /// This only makes sense in the context of a UnionAnalysis:
            /// it is an index into the parent Union's collection of generic parameters.
            RequiredGenerics : int list option
        }

    type CataUnionRecordField = (Ident * CataUnionBasicField) list

    [<RequireQualifiedAccess>]
    type CataUnionField =
        | Record of CataUnionRecordField
        | Basic of CataUnionBasicField

    /// Everything we'll need to know about a single union case within the
    /// user-provided DU.
    type RenderedUnionCase =
        {
            /// The name of the case within the `Instruction` state-machine DU
            /// which indicates "all the recursive calls are now resolved; you may proceed
            /// to pull recursive results from the stack and execute the cata directly"
            InstructionName : Ident
            /// This user-provided DU case
            CaseName : SynIdent
            /// The fields of this user-provided DU
            Fields : CataUnionField list
            /// The corresponding method of the appropriate cata, fully-qualified as a call
            /// into some specific cata
            CataMethodName : LongIdent
            /// The identifier of the method of the appropriate cata
            CataMethodIdent : SynIdent
            /// The Instruction case which instructs the state machine to pull anything
            /// necessary from the stacks and call into the cata.
            AssociatedInstruction : LongIdent
            /// Matching on an element of this union type, if you match against this
            /// left-hand side (and give appropriate field arguments), you will enter this union case.
            Match : LongIdent
        }

        member this.FlattenedFields : CataUnionBasicField list =
            this.Fields
            |> List.collect (fun f ->
                match f with
                | CataUnionField.Basic x -> [ x ]
                | CataUnionField.Record r -> r |> List.map snd
            )

    /// For a single user-provided DU (which is possibly one of several within a
    /// recursive knot), this is everything we need to know about it for the cata.
    type UnionAnalysis =
        {
            Accessibility : SynAccess option
            Typars : SynTyparDecl list
            /// The name of the stack we'll use for the results
            /// of returning from a descent into this union type,
            /// when performing the cata
            StackName : Ident
            /// The cases of this DU
            UnionCases : RenderedUnionCase list
            /// The Process instruction case which contains one of this union type.
            /// For example, the very first instruction processed will be one of these
            /// (i.e. when we enter the loop for the first time).
            /// The state machine interprets this instruction as "break me apart and
            /// descend recursively if necessary before coming back to me".
            AssociatedProcessInstruction : LongIdent
            /// Name of the parent type: e.g. in `type Foo = | Blah`, this is `Foo`.
            ParentTypeName : LongIdent
            /// The name of the generic type parameter we'll use within the cata
            /// to represent the result of cata'ing on this type.
            GenericName : Ident
            /// The name of the Cata type which represents "operate on this union case".
            CataTypeName : Ident
        }

    /// Returns a function:
    /// let run{Case} (cata : {cataName}<{typars}>) (x : {Case}) : {TyPar} =
    ///     let instructions = ResizeArray ()
    ///     instructions.Add (Instruction.Process{Case} e)
    ///     let {typar1}Results, {typar2}Results, ... = loop cata instructions
    ///     { for all non-relevant typars: }
    ///     if {typar}Results.Count > 0 then failwith "logic error"
    ///     Seq.exactlyOne {relevantTypar}Stack
    let createRunFunction
        (cataName : Ident)
        (userProvidedTypars : SynTyparDecl list)
        (allArtificialTypars : SynType list)
        (relevantTypar : SynType)
        (analysis : UnionAnalysis)
        : SynBinding
        =
        let relevantTypeName = analysis.ParentTypeName

        let allArtificialTyparNames =
            allArtificialTypars
            |> List.map (fun ty ->
                match ty with
                | SynType.Var (SynTypar.SynTypar (ident = ident), _) -> ident
                | _ -> failwith "logic error in generator"
            )

        let userProvidedTyparsForCase =
            analysis.Typars
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, ty)) -> SynType.var ty)

        let userProvidedTyparsForCata =
            userProvidedTypars
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, ty)) -> SynType.var ty)

        let relevantTyparName =
            match relevantTypar with
            | SynType.Var (SynTypar.SynTypar (ident = ident), _) -> ident
            | _ -> failwith "logic error in generator"

        let inputObjectType =
            let baseType = SynType.createLongIdent relevantTypeName

            if userProvidedTypars.Length = 0 then
                baseType
            else
                SynType.app' baseType userProvidedTyparsForCase

        // The object on which we'll run the cata
        let inputObject = SynPat.named "x" |> SynPat.annotateType inputObjectType

        let cataObject =
            SynPat.named "cata"
            |> SynPat.annotateType (
                SynType.app' (SynType.createLongIdent [ cataName ]) (userProvidedTyparsForCata @ allArtificialTypars)
            )

        [
            SynExpr.createLongIdent' analysis.AssociatedProcessInstruction
            |> SynExpr.applyTo (SynExpr.createLongIdent [ "x" ])
            |> SynExpr.paren
            |> SynExpr.applyFunction (SynExpr.createLongIdent [ "instructions" ; "Add" ])

            // TODO: add the "all other stacks are empty" sanity checks
            SynExpr.createIdent' (Ident.create (relevantTyparName.idText + "Stack") |> Ident.lowerFirstLetter)
            |> SynExpr.applyFunction (SynExpr.createLongIdent [ "Seq" ; "exactlyOne" ])
            |> SynExpr.createLet
                [
                    SynBinding.basicTuple
                        (allArtificialTyparNames
                         |> List.map (fun (t : Ident) ->
                             SynPat.namedI (Ident.create (t.idText + "Stack") |> Ident.lowerFirstLetter)
                         ))
                        (SynExpr.applyFunction
                            (SynExpr.applyFunction (SynExpr.createIdent "loop") (SynExpr.createIdent "cata"))
                            (SynExpr.createIdent "instructions"))
                ]
        ]
        |> SynExpr.sequential
        |> SynExpr.createLet
            [
                SynExpr.createIdent "ResizeArray"
                |> SynExpr.applyTo (SynExpr.CreateConst ())
                |> SynBinding.basic [ Ident.create "instructions" ] []
            ]
        |> SynExpr.typeAnnotate relevantTypar
        |> SynBinding.basic [ Ident.create ("run" + List.last(relevantTypeName).idText) ] [ cataObject ; inputObject ]
        |> SynBinding.withReturnAnnotation relevantTypar
        |> SynBinding.withXmlDoc (PreXmlDoc.create "Execute the catamorphism.")

    let getName (ty : SynTypeDefn) : LongIdent =
        match ty with
        | SynTypeDefn.SynTypeDefn (SynComponentInfo.SynComponentInfo (_, _, _, id, _, _, _, _), _, _, _, _, _) -> id

    let rec getNameUnion (unionType : SynType) : LongIdent option =
        match unionType with
        | SynType.LongIdent (SynLongIdent.SynLongIdent (name, _, _)) -> Some name
        | SynType.App (ty, _, _, _, _, _, _) -> getNameUnion ty
        | _ -> None

    let getNameKey (ty : SynTypeDefn) : string =
        getName ty |> List.map _.idText |> String.concat "/"

    // TODO: get rid of this function; it's causing some very spooky coupling at a distance
    let getNameKeyUnion (unionType : SynType) : string =
        match unionType with
        | SynType.LongIdent (SynLongIdent.SynLongIdent (name, _, _)) -> name |> List.map _.idText |> String.concat "/"
        | _ -> failwithf "unrecognised type: %+A" unionType

    /// Get the fields of this particular union case, and describe their relation to the
    /// recursive knot of user-provided DUs for which we are creating a cata.
    let analyse
        (availableGenerics : SynTyparDecl list)
        (allRecordTypes : SynTypeDefn list)
        (allUnionTypes : SynTypeDefn list)
        (argIndex : int)
        (fields : AdtNode list)
        : CataUnionBasicField list
        =
        let availableGenerics =
            availableGenerics
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident)

        let rec go (prefix : string) (name : Ident option) (ty : SynType) : CataUnionBasicField =
            let dealWithPrimitive (typeArgs : int list option) (ty : SynType) (typeName : LongIdent) =
                let key = typeName |> List.map _.idText |> String.concat "/"

                let isKnownUnion =
                    allUnionTypes |> List.exists (fun unionTy -> getNameKey unionTy = key)

                let knownRecord =
                    allRecordTypes
                    |> List.tryPick (fun recordTy -> if getNameKey recordTy = key then Some recordTy else None)

                if isKnownUnion then
                    {
                        FieldName = name
                        ArgName =
                            match name with
                            | Some n -> Ident.lowerFirstLetter n
                            | None -> Ident.create $"arg%s{prefix}"
                        Description = FieldDescription.Self ty
                        RequiredGenerics = typeArgs
                    }
                else
                    {
                        FieldName = name
                        ArgName =
                            match name with
                            | Some n -> Ident.lowerFirstLetter n
                            | None -> Ident.create $"arg%s{prefix}"
                        Description = FieldDescription.NonRecursive ty
                        RequiredGenerics = typeArgs
                    }

            let rec dealWithType (typeArgs : int list option) (stripped : SynType) =
                match stripped with
                | ListType child ->
                    let gone = go (prefix + "_") None child

                    match gone.Description with
                    | FieldDescription.NonRecursive ty ->
                        // Great, no recursion, just treat it as atomic
                        {
                            FieldName = name
                            ArgName =
                                match name with
                                | Some n -> Ident.lowerFirstLetter n
                                | None -> Ident.create $"arg%s{prefix}"
                            Description = FieldDescription.NonRecursive stripped
                            RequiredGenerics = typeArgs
                        }
                    | FieldDescription.Self ty ->
                        {
                            FieldName = name
                            ArgName =
                                match name with
                                | Some n -> Ident.lowerFirstLetter n
                                | None -> Ident.create $"arg%s{prefix}"
                            Description = FieldDescription.ListSelf ty
                            RequiredGenerics = typeArgs
                        }
                    | FieldDescription.ListSelf _ -> failwith "Deeply nested lists not currently supported"
                | PrimitiveType _ ->
                    {
                        FieldName = name
                        ArgName =
                            match name with
                            | Some n -> Ident.lowerFirstLetter n
                            | None -> Ident.create $"arg%s{prefix}"
                        Description = FieldDescription.NonRecursive stripped
                        RequiredGenerics = typeArgs
                    }
                | SynType.App (ty, _, childTypeArgs, _, _, _, _) ->
                    match typeArgs with
                    | Some _ -> failwithf "Nested applications of types not supported in %+A" ty
                    | None ->
                        let childTypeArgs =
                            childTypeArgs
                            |> List.map (fun generic ->
                                let generic =
                                    match generic with
                                    | SynType.Var (SynTypar.SynTypar (name, _, _), _) -> name
                                    | _ -> failwithf "Unrecognised generic arg: %+A" generic

                                availableGenerics
                                |> List.findIndex (fun knownGeneric -> knownGeneric.idText = generic.idText)
                            )

                        dealWithType (Some childTypeArgs) (SynType.stripOptionalParen ty)
                | SynType.LongIdent (SynLongIdent.SynLongIdent (ty, _, _)) -> dealWithPrimitive typeArgs stripped ty
                | SynType.Var (typar, _) ->
                    {
                        FieldName = name
                        ArgName =
                            match name with
                            | Some n -> Ident.lowerFirstLetter n
                            | None -> Ident.create $"arg%s{prefix}"
                        Description = FieldDescription.NonRecursive ty
                        RequiredGenerics = typeArgs
                    }

                | _ -> failwithf "Unrecognised type: %+A" stripped

            let stripped = SynType.stripOptionalParen ty
            dealWithType None stripped

        fields |> List.mapi (fun i x -> go $"%i{argIndex}_%i{i}" x.Name x.Type)

    /// Returns whether this type recursively contains a Self, and the type which
    /// the Instruction case is going to have to store to obtain this field.
    /// (For example, a `self list` will need to store an int, namely the number
    /// of recursive results to pull from the stack just before we feed them
    /// into the cata.)
    let rec toInstructionCase (field : FieldDescription) : bool * SynType option =
        match field with
        | FieldDescription.NonRecursive ty -> false, Some ty
        | FieldDescription.Self ty -> true, None
        | FieldDescription.ListSelf ty ->
            // store the length of the list
            true, Some SynType.int

    type InstructionCase =
        {
            Name : Ident
            Fields : AdtNode list
        }

    let getInstructionCaseName (thisUnionType : SynTypeDefn) (caseName : SynIdent) : Ident =
        match caseName with
        | SynIdent.SynIdent (ident, _) ->
            (List.last (getName thisUnionType)).idText + "_" + ident.idText |> Ident.create

    /// Given the input `| Pair of Expr * Expr * PairOpKind`,
    /// strips out any members which contain recursive calls.
    /// Stores a list as an int which is "the length of the list".
    /// TODO: support other compound types.
    let getRecursiveInstruction (case : RenderedUnionCase) : InstructionCase option =
        let hasRecursion, cases =
            ((false, []), case.FlattenedFields)
            ||> List.fold (fun (hasRecursion, cases) field ->
                let newHasRecursion, case = toInstructionCase field.Description

                let cases =
                    match case with
                    | None -> cases
                    | Some case -> (field.FieldName, case) :: cases

                hasRecursion || newHasRecursion, cases
            )

        if not hasRecursion then
            // No recursive instructions required; we'll be feeding the data
            // straight into the cata without any stack manipulation.
            None
        else

        let fields =
            cases
            |> List.rev
            |> List.map (fun (name, ty) ->
                {
                    Name = name |> Option.map Ident.lowerFirstLetter
                    Type = ty
                    // TODO this is definitely wrong
                    GenericsOfParent = []
                }
            )

        {
            Name = case.InstructionName
            Fields = fields
        }
        |> Some

    /// The instruction to "process an Expr"; the loop will have to descend
    /// into this Expr and break it down to discover what recursive calls
    /// and calls to the cata this will imply making.
    let baseCases (unions : UnionAnalysis list) : InstructionCase list =
        unions
        |> List.map (fun union ->
            {
                Name = List.last union.AssociatedProcessInstruction
                Fields =
                    {
                        Name = None
                        Type =
                            let name = SynType.createLongIdent union.ParentTypeName

                            match union.Typars with
                            | [] -> name
                            | typars ->
                                let typars = typars |> List.map (fun (SynTyparDecl (_, typar)) -> SynType.var typar)

                                SynType.app' name typars

                        GenericsOfParent = union.Typars
                    }
                    |> List.singleton
            }
        )

    /// The instruction to "pull recursive results from the stack, and then call into the cata".
    let recursiveCases (allUnionTypes : UnionAnalysis list) : InstructionCase list =
        allUnionTypes
        |> List.collect (fun union -> union.UnionCases |> List.choose getRecursiveInstruction)

    /// Build the DU which defines the states our state machine can be in.
    let createInstructionType (analysis : UnionAnalysis list) : SynTypeDefn =
        let parentGenerics =
            analysis
            |> List.collect _.Typars
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
            |> List.distinct
            |> List.map (fun i ->
                SynTyparDecl.SynTyparDecl ([], SynTypar.SynTypar (Ident.create i, TyparStaticReq.None, false))
            )

        // One union case for each union type, and then
        // a union case for each union case which contains a recursive reference.
        let casesFromProcess : SynUnionCase list =
            baseCases analysis
            |> List.map (fun unionCase ->
                let fields =
                    unionCase.Fields
                    |> List.map (fun field ->
                        // TODO: adjust type parameters
                        {
                            SynFieldData.Type = field.Type
                            Attrs = []
                            Ident = field.Name
                        }
                    )

                {
                    Name = unionCase.Name
                    XmlDoc = None
                    Access = None
                    Attributes = []
                    Fields = fields
                }
                |> SynUnionCase.create
            )

        let casesFromCases =
            recursiveCases analysis
            |> List.map (fun case ->
                {
                    UnionCase.Name = case.Name
                    XmlDoc = None
                    Access = None
                    Attributes = []
                    Fields =
                        case.Fields
                        |> List.map (fun field ->
                            {
                                SynFieldData.Type = field.Type
                                Attrs = []
                                Ident = field.Name
                            }
                        )
                }
                |> SynUnionCase.create
            )

        let cases = casesFromProcess @ casesFromCases

        let typars =
            if analysis |> List.forall (fun x -> x.Typars.IsEmpty) then
                []
            else
                analysis
                |> List.collect _.Typars
                |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
                |> List.distinct
                |> List.map (fun i ->
                    SynTyparDecl.SynTyparDecl ([], SynTypar.SynTypar (Ident.create i, TyparStaticReq.None, false))
                )

        SynTypeDefnRepr.union cases
        |> SynTypeDefn.create (
            SynComponentInfo.create (Ident.create "Instruction")
            |> SynComponentInfo.withGenerics typars
            |> SynComponentInfo.withAccessibility (SynAccess.Private range0)
            |> SynComponentInfo.addAttributes [ SynAttribute.requireQualifiedAccess ]
        )

    /// Build the cata interfaces, which a user will instantiate to specify a particular
    /// catamorphism. This produces one interface per input union type.
    ///
    /// Say that CreateCatamorphism-tagged types form the set T.
    /// Assert that each U in T is a discriminated union.
    /// For each type U in T, assign a generic parameter 'ret<U>.
    /// For each U:
    ///   * Define the type [U]Cata, generic on all the parameters {'ret<U> : U in T}.
    ///   * For each DU case C in type U:
    ///     * create a method in [U]Cata, whose return value is 'ret<U> and whose args are the fields of the case C
    ///     * any occurrence in a field of an input value of type equal to any element of T (say type V) is replaced by 'ret<V>
    let createCataStructure (analyses : UnionAnalysis list) : SynTypeDefn list =
        // Obtain the generic parameter for a UnionAnalysis by dotting into this
        // with `case.GenericName.idText`.
        // Remember that this is essentially unordered!
        let generics =
            analyses
            |> List.map (fun case ->
                case.GenericName.idText, SynTypar.SynTypar (case.GenericName, TyparStaticReq.None, false)
            )
            |> Map.ofList

        let orderedGenerics =
            analyses
            |> List.map (fun case -> SynTyparDecl.SynTyparDecl ([], generics.[case.GenericName.idText]))

        analyses
        |> List.map (fun analysis ->
            let componentInfo =
                SynComponentInfo.create analysis.CataTypeName
                // TODO: better docstring
                |> SynComponentInfo.withDocString (
                    PreXmlDoc.create "Description of how to combine cases during a fold"
                )
                |> SynComponentInfo.withGenerics (analysis.Typars @ orderedGenerics)

            analysis.UnionCases
            |> List.map (fun case ->
                let arity =
                    SynValInfo.SynValInfo (
                        case.Fields |> List.map (fun field -> [ SynArgInfo.empty ]),
                        SynArgInfo.empty
                    )

                (SynType.var generics.[analysis.GenericName.idText], List.rev case.FlattenedFields)
                ||> List.fold (fun acc field ->
                    let place : SynType =
                        match field.Description with
                        | FieldDescription.Self ty -> SynType.var generics.[getNameKeyUnion ty]
                        | FieldDescription.ListSelf ty ->
                            SynType.var generics.[getNameKeyUnion ty] |> SynType.appPostfix "list"
                        | FieldDescription.NonRecursive ty ->
                            match field.RequiredGenerics with
                            | None -> ty
                            | Some generics ->
                                generics
                                |> List.map (fun i ->
                                    let (SynTyparDecl.SynTyparDecl (_, typar)) = analysis.Typars.[i]
                                    SynType.var typar
                                )
                                |> SynType.app' ty

                    let domain =
                        field.FieldName
                        |> Option.map Ident.lowerFirstLetter
                        |> SynType.signatureParamOfType [] place false

                    acc |> SynType.funFromDomain domain
                )
                |> SynMemberDefn.abstractMember
                    []
                    case.CataMethodIdent
                    None
                    arity
                    (PreXmlDoc.create $"How to operate on the %s{List.last(case.Match).idText} case")
            )
            |> SynTypeDefnRepr.interfaceType
            |> SynTypeDefn.create componentInfo
        )

    /// Build a record which contains one of every cata type.
    /// That is, define a type Cata<{'ret<U> for U in T}>
    /// with one member for each U, namely of type [U]Cata<{'ret<U> for U in T}>.
    let createCataRecord (cataName : Ident) (doc : PreXmlDoc) (analysis : UnionAnalysis list) : SynTypeDefn =
        // An artificial generic for each union type
        let generics =
            analysis
            |> List.map (fun analysis -> SynTypar.SynTypar (analysis.GenericName, TyparStaticReq.None, false))

        // A field for each cata
        let fields =
            analysis
            |> List.map (fun analysis ->
                let nameForDoc = List.last(analysis.ParentTypeName).idText

                let doc =
                    PreXmlDoc.create $"How to perform a fold (catamorphism) over the type %s{nameForDoc}"

                let artificialGenerics = generics |> List.map SynType.var

                let userInputGenerics =
                    analysis.Typars
                    |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
                    |> List.distinct
                    |> List.map (fun i -> SynType.var (SynTypar.SynTypar (Ident.create i, TyparStaticReq.None, false)))

                let ty =
                    SynType.app'
                        (SynType.createLongIdent [ analysis.CataTypeName ])
                        (userInputGenerics @ artificialGenerics)

                SynField.SynField (
                    [],
                    false,
                    Some (List.last analysis.ParentTypeName),
                    ty,
                    false,
                    doc,
                    None,
                    range0,
                    {
                        LeadingKeyword = None
                    }
                )
            )

        // A "real" generic for each generic in the user-provided type
        let genericsFromUserInput =
            analysis
            |> List.collect _.Typars
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
            |> List.distinct
            |> List.map (fun i ->
                SynTyparDecl.SynTyparDecl ([], SynTypar.SynTypar (Ident.create i, TyparStaticReq.None, false))
            )

        let genericsFromCata =
            generics |> List.map (fun ty -> SynTyparDecl.SynTyparDecl ([], ty))

        let componentInfo =
            SynComponentInfo.create cataName
            |> SynComponentInfo.withGenerics (genericsFromUserInput @ genericsFromCata)
            |> SynComponentInfo.withDocString doc

        SynTypeDefnRepr.record fields |> SynTypeDefn.create componentInfo

    let makeUnionAnalyses
        (cataVarName : Ident)
        (allRecordTypes : SynTypeDefn list)
        (allUnionTypes : SynTypeDefn list)
        : UnionAnalysis list
        =
        let recordTypes =
            allRecordTypes
            |> List.map (fun ty -> List.last(getName ty).idText, AstHelper.getRecordFields ty)
            |> Map.ofList

        allUnionTypes
        |> List.map (fun unionType ->
            let cases, typars, access = AstHelper.getUnionCases unionType

            let cases =
                cases
                |> List.map (fun prod ->
                    let fields =
                        prod.Fields
                        |> List.indexed
                        |> List.collect (fun (i, node) ->
                            match getNameUnion node.Type with
                            | None ->
                                analyse typars allRecordTypes allUnionTypes i [ node ]
                                |> List.map CataUnionField.Basic
                            | Some name ->

                            match Map.tryFind (List.last(name).idText) recordTypes with
                            | None ->
                                analyse typars allRecordTypes allUnionTypes i [ node ]
                                |> List.map CataUnionField.Basic
                            | Some fields ->
                                List.zip fields (analyse typars allRecordTypes allUnionTypes i fields)
                                |> List.map (fun (field, analysis) -> Option.get field.Name, analysis)
                                |> CataUnionField.Record
                                |> List.singleton
                        )

                    prod.Name, fields
                )

            let unionTypeName = getName unionType

            {
                Typars = typars
                Accessibility = access
                StackName =
                    List.last(getName unionType).idText + "Stack"
                    |> Ident.create
                    |> Ident.lowerFirstLetter
                UnionCases =
                    cases
                    |> List.map (fun (name, analysis) ->
                        let instructionName = getInstructionCaseName unionType name

                        let unionCaseName =
                            match name with
                            | SynIdent (ident, _) -> ident

                        {
                            InstructionName = instructionName
                            Fields = analysis
                            CaseName = name
                            CataMethodName = cataVarName :: unionTypeName @ [ unionCaseName ]
                            CataMethodIdent = SynIdent.SynIdent (unionCaseName, None)
                            AssociatedInstruction = [ Ident.create "Instruction" ; instructionName ]
                            Match = unionTypeName @ [ unionCaseName ]
                        }
                    )
                AssociatedProcessInstruction =
                    [
                        "Instruction"
                        // such jank!
                        "Process__" + List.last(unionTypeName).idText
                    ]
                    |> List.map Ident.create
                ParentTypeName = getName unionType
                GenericName = getName unionType |> List.map _.idText |> String.concat "" |> Ident.create
                CataTypeName = List.last(getName unionType).idText + "CataCase" |> Ident.create
            }
        )

    let callCataAndPushResult (resultStackName : Ident) (unionCase : RenderedUnionCase) : SynExpr =
        (SynExpr.createLongIdent' unionCase.CataMethodName, unionCase.FlattenedFields)
        ||> List.fold (fun body caseDesc -> SynExpr.applyFunction body (SynExpr.createIdent' caseDesc.ArgName))
        |> SynExpr.pipeThroughFunction (SynExpr.createLongIdent' (resultStackName :: [ Ident.create "Add" ]))

    /// Create the state-machine matches which deal with receiving the instruction
    /// to "process one of the user-specified DU cases, pushing recursion instructions onto
    /// the instruction stack".
    /// It very rarely involves invoking the cata; that happens only if there's no recursion.
    let createBaseMatchClause (analysis : UnionAnalysis) : SynMatchClause =
        let matchCases =
            analysis.UnionCases
            |> List.map (fun unionCase ->
                let name =
                    match unionCase.CaseName with
                    | SynIdent (ident, _) -> ident

                let _, nonRecursiveArgs, selfArgs, listSelfArgs =
                    ((0, [], [], []), unionCase.FlattenedFields)
                    ||> List.fold (fun (i, nonRec, self, listSelf) caseDesc ->
                        match caseDesc.Description with
                        | FieldDescription.NonRecursive ty ->
                            i + 1, (i, caseDesc.ArgName, ty) :: nonRec, self, listSelf
                        | FieldDescription.Self ty -> i + 1, nonRec, (i, caseDesc.ArgName, ty) :: self, listSelf
                        | FieldDescription.ListSelf ty -> i + 1, nonRec, self, (i, caseDesc.ArgName, ty) :: listSelf
                    )

                let matchBody =
                    if nonRecursiveArgs.Length = unionCase.FlattenedFields.Length then
                        // directly call the cata
                        callCataAndPushResult analysis.StackName unionCase
                    else
                    // There's a recursive type in here, so we'll have to make some calls
                    // and then come back.

                    // The instruction to process us again once our inputs are ready:
                    let reprocessCommand =
                        if selfArgs.Length = unionCase.FlattenedFields.Length then
                            SynExpr.createLongIdent' unionCase.AssociatedInstruction
                        else
                            // We need to tell ourselves each non-rec arg, and the length of each input list.
                            listSelfArgs
                            |> List.map (fun (i, argName, _) ->
                                i,
                                SynExpr.paren (
                                    SynExpr.applyFunction
                                        (SynExpr.createLongIdent [ "List" ; "length" ])
                                        (SynExpr.createIdent' argName)
                                )
                            )
                            |> List.append (
                                nonRecursiveArgs |> List.map (fun (i, arg, _) -> i, SynExpr.createIdent' arg)
                            )
                            |> List.sortBy fst
                            |> List.map snd
                            |> SynExpr.tuple
                            |> SynExpr.applyFunction (SynExpr.createLongIdent' unionCase.AssociatedInstruction)
                            |> SynExpr.paren
                        |> SynExpr.applyFunction (SynExpr.createLongIdent [ "instructions" ; "Add" ])

                    [
                        yield reprocessCommand

                        for i, caseDesc in Seq.indexed unionCase.FlattenedFields do
                            match caseDesc.Description with
                            | NonRecursive synType ->
                                // Nothing to do, because we're not calling the cata yet
                                ()
                            | ListSelf synType ->
                                // Tell our future self to process the list elements first.
                                yield
                                    SynExpr.ForEach (
                                        DebugPointAtFor.Yes range0,
                                        DebugPointAtInOrTo.Yes range0,
                                        SeqExprOnly.SeqExprOnly false,
                                        true,
                                        SynPat.named "elt",
                                        SynExpr.createIdent' caseDesc.ArgName,
                                        SynExpr.applyFunction
                                            (SynExpr.createLongIdent [ "instructions" ; "Add" ])
                                            (SynExpr.paren (
                                                SynExpr.applyFunction
                                                    (SynExpr.createLongIdent' analysis.AssociatedProcessInstruction)
                                                    (SynExpr.createIdent "elt")
                                            )),
                                        range0
                                    )
                            | Self synType ->
                                // And push the instruction to process each recursive call
                                // onto the stack.
                                yield
                                    // TODO: use an AssociatedProcessInstruction instead
                                    SynExpr.createLongIdent
                                        [
                                            "Instruction"
                                            // TODO wonky domain
                                            "Process" + "__" + List.last(getNameUnion(synType).Value).idText
                                        ]
                                    |> SynExpr.applyTo (SynExpr.createIdent' caseDesc.ArgName)
                                    |> SynExpr.paren
                                    |> SynExpr.applyFunction (SynExpr.createLongIdent [ "instructions" ; "Add" ])
                    ]
                    |> SynExpr.sequential

                let matchLhs =
                    if not unionCase.Fields.IsEmpty then
                        unionCase.Fields
                        |> List.mapi (fun i case ->
                            match case with
                            | CataUnionField.Basic case -> SynPat.namedI (Ident.lowerFirstLetter case.ArgName)
                            | CataUnionField.Record fields ->
                                let fields =
                                    fields
                                    |> List.map (fun (name, field) ->
                                        ([], name), range0, SynPat.namedI (Ident.lowerFirstLetter name)
                                    )

                                SynPat.Record (fields, range0)
                        )
                        |> SynPat.tuple
                        |> List.singleton
                    else
                        []

                SynMatchClause.create (SynPat.identWithArgs unionCase.Match (SynArgPats.create matchLhs)) matchBody
            )

        SynExpr.createMatch (SynExpr.createIdent "x") matchCases
        |> SynMatchClause.create (
            SynPat.identWithArgs analysis.AssociatedProcessInstruction (SynArgPats.createNamed [ "x" ])
        )

    /// Create the state-machine matches which deal with receiving the instruction
    /// to "pull recursive results from the result stacks, and invoke the cata".
    let createRecursiveMatchClauses (analyses : UnionAnalysis list) : SynMatchClause list =
        let inputStacks =
            analyses
            |> Seq.map (fun a ->
                // TODO this is jank
                (List.last a.ParentTypeName).idText, a.StackName
            )
            |> Map.ofSeq

        analyses
        |> List.collect (fun analysis ->
            analysis.UnionCases
            |> List.choose (fun unionCase ->
                // We already know there is a recursive reference somewhere
                // in `analysis`.
                if
                    unionCase.FlattenedFields
                    |> List.exists (fun case ->
                        match case.Description with
                        | NonRecursive _ -> false
                        | _ -> true
                    )
                then
                    Some unionCase
                else
                    None
            )
            |> List.map (fun unionCase ->
                let pat =
                    unionCase.FlattenedFields
                    |> Seq.mapi (fun i x -> (i, x))
                    |> Seq.choose (fun (i, case) ->
                        match case.Description with
                        | FieldDescription.NonRecursive _ -> case.ArgName |> SynPat.namedI |> Some
                        | FieldDescription.ListSelf _ -> case.ArgName |> SynPat.namedI |> Some
                        | FieldDescription.Self _ -> None
                    )
                    |> Seq.toList
                    |> SynArgPats.create
                    |> SynPat.identWithArgs unionCase.AssociatedInstruction

                let populateArgs =
                    unionCase.FlattenedFields
                    |> List.choose (fun field ->
                        match field.Description with
                        | NonRecursive _ ->
                            // this was passed in already in the match
                            None
                        | Self synType ->
                            // pull the one entry from the stack
                            // let {field.ArgName} = {appropriateStack}.[SynExpr.minusN {appropriateStack.Count} 1]
                            // {appropriateStack}.RemoveRange (SynExpr.minusN {appropriateStack.Count} 1)
                            // TODO: this is jank
                            let stackName = inputStacks.[List.last(getNameUnion(synType).Value).idText]

                            SynExpr.minusN (SynLongIdent.create [ stackName ; Ident.create "Count" ]) 1
                            |> SynExpr.paren
                            |> SynExpr.applyFunction (
                                SynExpr.createLongIdent' [ stackName ; Ident.create "RemoveAt" ]
                            )
                            |> SynExpr.createLet
                                [
                                    SynExpr.DotIndexedGet (
                                        SynExpr.createIdent' stackName,
                                        SynExpr.minusN (SynLongIdent.create [ stackName ; Ident.create "Count" ]) 1,
                                        range0,
                                        range0
                                    )
                                    |> SynBinding.basic [ field.ArgName ] []
                                ]
                            |> Some
                        | ListSelf synType ->
                            // TODO: also jank
                            let stackName = inputStacks.[List.last(getNameUnion(synType).Value).idText]

                            let vals =
                                SynExpr.ComputationExpr (
                                    false,
                                    SynExpr.For (
                                        DebugPointAtFor.Yes range0,
                                        DebugPointAtInOrTo.Yes range0,
                                        Ident.create "i",
                                        Some range0,
                                        SynExpr.minusN (SynLongIdent.create [ stackName ; Ident.create "Count" ]) 1,
                                        false,
                                        SynExpr.minus
                                            (SynLongIdent.create [ stackName ; Ident.create "Count" ])
                                            (SynExpr.createIdent' field.ArgName),
                                        SynExpr.YieldOrReturn (
                                            (true, false),
                                            SynExpr.DotIndexedGet (
                                                SynExpr.createIdent' stackName,
                                                SynExpr.createIdent "i",
                                                range0,
                                                range0
                                            ),
                                            range0
                                        ),
                                        range0
                                    ),
                                    range0
                                )
                                |> SynExpr.applyFunction (SynExpr.createIdent "seq")
                                |> SynExpr.pipeThroughFunction (SynExpr.createLongIdent [ "Seq" ; "toList" ])
                                |> SynBinding.basic [ field.ArgName ] []

                            let shadowedIdent = Ident.create (field.ArgName.idText + "_len")

                            [
                                SynExpr.minus
                                    (SynLongIdent.create [ stackName ; Ident.create "Count" ])
                                    (SynExpr.createIdent' shadowedIdent)
                                SynExpr.createIdent' shadowedIdent
                            ]
                            |> SynExpr.tuple
                            |> SynExpr.applyFunction (
                                SynExpr.createLongIdent' [ stackName ; Ident.create "RemoveRange" ]
                            )
                            |> SynExpr.createLet [ vals ]
                            |> SynExpr.createLet
                                [ SynBinding.basic [ shadowedIdent ] [] (SynExpr.createIdent' field.ArgName) ]
                            |> Some
                    )

                (populateArgs @ [ callCataAndPushResult analysis.StackName unionCase ])
                |> SynExpr.sequential
                |> SynMatchClause.create pat
            )
        )

    let createLoopFunction (cataTypeName : Ident) (cataVarName : Ident) (analysis : UnionAnalysis list) : SynBinding =
        let userSuppliedGenerics =
            analysis
            |> List.collect _.Typars
            |> List.map (fun (SynTyparDecl.SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
            |> List.distinct
            |> List.map (fun i -> SynTypar.SynTypar (Ident.create i, TyparStaticReq.None, false))

        let instructionsArrType =
            if not userSuppliedGenerics.IsEmpty then
                userSuppliedGenerics |> List.map SynType.var |> SynType.app "Instruction"
            else
                SynType.named "Instruction"

        let cataGenerics =
            [
                for generic in userSuppliedGenerics do
                    yield SynType.var generic
                for case in analysis do
                    yield SynType.var (SynTypar.SynTypar (case.GenericName, TyparStaticReq.None, false))
            ]

        let args =
            [
                SynPat.namedI cataVarName
                |> SynPat.annotateType (SynType.app' (SynType.createLongIdent [ cataTypeName ]) cataGenerics)

                SynPat.named "instructions"
                |> SynPat.annotateType (SynType.app "ResizeArray" [ instructionsArrType ])
            ]

        let baseMatchClauses = analysis |> List.map createBaseMatchClause

        let recMatchClauses = createRecursiveMatchClauses analysis

        let matchStatement =
            SynExpr.createMatch (SynExpr.createIdent "currentInstruction") (baseMatchClauses @ recMatchClauses)

        let body =
            [
                SynExpr.applyFunction
                    (SynExpr.createLongIdent [ "instructions" ; "RemoveAt" ])
                    (SynExpr.paren (SynExpr.minusN (SynLongIdent.createS' [ "instructions" ; "Count" ]) 1))
                matchStatement
            ]
            |> SynExpr.sequential
            |> SynExpr.createLet
                [
                    SynExpr.DotIndexedGet (
                        SynExpr.createIdent "instructions",
                        SynExpr.minusN (SynLongIdent.createS' [ "instructions" ; "Count" ]) 1,
                        range0,
                        range0
                    )
                    |> SynBinding.basic [ Ident.create "currentInstruction" ] []
                ]

        let body =
            SynExpr.sequential
                [
                    SynExpr.createWhile
                        (SynExpr.greaterThan
                            (SynExpr.CreateConst 0)
                            (SynExpr.createLongIdent [ "instructions" ; "Count" ]))
                        body
                    SynExpr.tupleNoParen (
                        analysis
                        |> List.map (fun unionAnalysis -> [ unionAnalysis.StackName ] |> SynExpr.createLongIdent')
                    )
                ]

        let body =
            (body, analysis)
            ||> List.fold (fun body unionCase ->
                body
                |> SynExpr.createLet
                    [
                        (SynExpr.createIdent "ResizeArray")
                        |> SynExpr.typeApp
                            [
                                SynType.var (SynTypar.SynTypar (unionCase.GenericName, TyparStaticReq.None, false))
                            ]
                        |> SynExpr.applyTo (SynExpr.CreateConst ())
                        |> SynBinding.basic [ unionCase.StackName ] []
                    ]
            )

        SynBinding.basic [ Ident.create "loop" ] args body
        |> SynBinding.withAccessibility (Some (SynAccess.Private range0))

    let createModule
        (opens : SynOpenDeclTarget list)
        (ns : LongIdent)
        (cataName : SynExpr, taggedType : SynTypeDefn)
        (allUnionTypes : SynTypeDefn list)
        (allRecordTypes : SynTypeDefn list)
        : SynModuleOrNamespace
        =
        let cataName =
            match cataName |> SynExpr.stripOptionalParen with
            | SynExpr.Const (SynConst.String (name, _, _), _) -> Ident.create name
            | _ -> failwith "Cata name in attribute must be literally a string, sorry"

        let parentName = List.last (getName taggedType) |> _.idText
        let moduleName = parentName + "Cata" |> Ident.create

        let modInfo =
            SynComponentInfo.create moduleName
            |> SynComponentInfo.withDocString (
                PreXmlDoc.create $"Methods to perform a catamorphism over the type %s{parentName}"
            )
            |> SynComponentInfo.addAttributes [ SynAttribute.requireQualifiedAccess ]

        let cataVarName = Ident.create "cata"
        let analysis = makeUnionAnalyses cataVarName allRecordTypes allUnionTypes

        let allTypars =
            allUnionTypes
            |> List.map (fun unionType ->
                List.last (getName unionType)
                |> fun x -> x.idText + "Ret"
                |> Ident.create
                |> fun x -> SynTypar.SynTypar (x, TyparStaticReq.None, false)
                |> SynType.var
            )

        let userProvidedGenerics =
            analysis
            |> List.collect _.Typars
            |> List.map (fun (SynTyparDecl (_, SynTypar.SynTypar (ident, _, _))) -> ident.idText)
            |> List.distinct
            |> List.map (fun x ->
                SynTyparDecl.SynTyparDecl ([], SynTypar.SynTypar (Ident.create x, TyparStaticReq.None, false))
            )

        let runFunctions =
            List.zip analysis allTypars
            |> List.map (fun (analysis, relevantTypar) ->
                createRunFunction cataName userProvidedGenerics allTypars relevantTypar analysis
            )

        let cataStructures =
            createCataStructure analysis
            |> List.map (fun repr -> SynModuleDecl.Types ([ repr ], range0))

        let loopFunction = createLoopFunction cataName cataVarName analysis

        let recordDoc =
            PreXmlDoc.create
                $"Specifies how to perform a fold (catamorphism) over the type %s{parentName} and its friends."

        let cataRecord =
            SynModuleDecl.Types ([ createCataRecord cataName recordDoc analysis ], range0)

        [
            for openStatement in opens do
                yield SynModuleDecl.openAny openStatement
            yield! cataStructures
            yield cataRecord
            yield
                [
                    SynModuleDecl.Types ([ createInstructionType analysis ], range0)
                    SynModuleDecl.createLets (loopFunction :: runFunctions)
                ]
                |> SynModuleDecl.nestedModule modInfo
        ]
        |> SynModuleOrNamespace.createNamespace ns

    /// For each namespace/module, grab the types which are defined in consecutive `and`-knots in that namespace/module,
    /// and also return the fully-qualified namespace/module name alongside that group of types.
    /// A given module LongIdent may show up many times in the output: once for each recursive knot.
    // Function originally inspired by https://github.com/MoiraeSoftware/myriad/blob/3c9818faabf9d508c10c28d5ecd26e66fafb48a1/src/Myriad.Core/Ast.fs#L160
    // but there's really only one reasonable implementation of this type signature and semantics.
    let groupedTypeDefns (ast : ParsedInput) : (LongIdent * SynTypeDefn list) list =
        let rec extractTypes (decls : SynModuleDecl list) (ns : LongIdent) =
            decls
            |> List.collect (fun moduleDecl ->
                match moduleDecl with
                | SynModuleDecl.Types (types, _) -> [ ns, types ]
                | SynModuleDecl.NestedModule (SynComponentInfo (_, _, _, longId, _, _, _, _), _, decls, _, _, _) ->
                    let combined = longId |> List.append ns
                    extractTypes decls combined
                | _ -> []
            )

        match ast with
        | ParsedInput.ImplFile (ParsedImplFileInput (_, _, _, _, _, contents, _, _, _)) ->
            contents
            |> List.collect (fun (SynModuleOrNamespace (namespaceId, _, _, moduleDecls, _, _, _, _, _)) ->
                extractTypes moduleDecls namespaceId
            )
        | _ -> []

open Myriad.Core

/// Myriad generator that provides a catamorphism for an algebraic data type.
[<MyriadGenerator("create-catamorphism")>]
type CreateCatamorphismGenerator () =

    interface IMyriadGenerator with
        member _.ValidInputExtensions = [ ".fs" ]

        member _.Generate (context : GeneratorContext) =
            let ast, _ =
                Ast.fromFilename context.InputFilename |> Async.RunSynchronously |> Array.head

            let types = CataGenerator.groupedTypeDefns ast

            let opens = AstHelper.extractOpens ast

            let namespaceAndTypes =
                types
                |> List.choose (fun (ns, types) ->
                    let typeWithAttr =
                        types
                        |> List.tryPick (fun ty ->
                            match SynTypeDefn.getAttribute typeof<CreateCatamorphismAttribute>.Name ty with
                            | None -> None
                            | Some attr -> Some (attr.ArgExpr, ty)
                        )

                    match typeWithAttr with
                    | Some taggedType ->
                        let unions, records, others =
                            (([], [], []), types)
                            ||> List.fold (fun
                                               (unions, records, others)
                                               (SynTypeDefn.SynTypeDefn (_, repr, _, _, _, _) as ty) ->
                                match repr with
                                | SynTypeDefnRepr.Simple (SynTypeDefnSimpleRepr.Union _, _) ->
                                    ty :: unions, records, others
                                | SynTypeDefnRepr.Simple (SynTypeDefnSimpleRepr.Record _, _) ->
                                    unions, ty :: records, others
                                | _ -> unions, records, ty :: others
                            )

                        if not others.IsEmpty then
                            failwith
                                $"Error: all types recursively defined together with a CreateCatamorphism type must be discriminated unions or records. %+A{others}"

                        Some (ns, taggedType, unions, records)
                    | _ -> None
                )

            let modules =
                namespaceAndTypes
                |> List.map (fun (ns, taggedType, unions, records) ->
                    CataGenerator.createModule opens ns taggedType unions records
                )

            Output.Ast modules