Dictionaries.ff

import Syntax
import Unification
import Environment

data Dictionaries(
    instances: Map[InstanceKey, InstanceValue]
)

make(modules: List[Module]): Dictionaries {
    Dictionaries(Unification.make(modules, False).instances)
}

fail[T](at: Location, message: String): T {
    panic(message + " " + at.show())
}

extend self: Dictionaries {

    processModule(module: Module, otherModules: List[Module]): Module {
        let environment = Environment.make(module, otherModules, alreadyFlat = True)
        let functionSignatures = environment.symbols.pairs().collect {
            | Pair(name, s) {!s.isVariable} => Some(Pair(name, s.signature))
            | _ => None
        }.toMap()

        let lets = module.lets.map { self.processLetDefinition(functionSignatures, _) }
        let functions = module.functions.map { self.processFunctionDefinition(functionSignatures, _) }
        let extends = module.extends.map { self.processExtendDefinition(functionSignatures, _) }
        //let traits = module.traits.map { self.processTraitDefinition(functionSignatures, _) }
        let instances = module.instances.map { self.processInstanceDefinition(functionSignatures, _) }
        module.Module(
            //traits = traits,
            instances = instances,
            extends = extends,
            lets = lets,
            functions = functions,
        )
    }

    processLetDefinition(functions: Map[String, Signature], definition: DLet): DLet {
        definition.DLet(
            value = self.processTerm(functions, definition.value)
        )
    }

    processExtendDefinition(functions: Map[String, Signature], definition: DExtend): DExtend {
        definition.DExtend(
            methods = definition.methods.map { self.processFunctionDefinition(functions, _) }
        )
    }

    processFunctionDefinition(functions: Map[String, Signature], definition: DFunction): DFunction {
        let instances = constraintsToInstances(definition.signature.constraints)
        let self2 = self.Dictionaries(instances = self.instances.addAll(instances))
        definition.DFunction(
            body = definition.body.mapFirefly {self2.processLambda(functions, _)}
        )
    }

    processInstanceDefinition(functions: Map[String, Signature], definition: DInstance): DInstance {
        let instances = constraintsToInstances(definition.constraints)
        let self2 = self.Dictionaries(instances = self.instances.addAll(instances))
        definition.DInstance(
            methods = definition.methods.map { self2.processFunctionDefinition(functions, _) }
        )
    }

    processLambda(functions: Map[String, Signature], definition: Lambda): Lambda {
        definition.Lambda(
            cases = definition.cases.map { case =>
                // TODO: Remove the case variables from the functions
                case.MatchCase(
                    guards = case.guards.map { g => g.MatchGuard(term = self.processTerm(functions, g.term)) }
                    body = self.processTerm(functions, case.body)
                )
            }
        )
    }

    processTerm(functions: Map[String, Signature], term: Term): Term {
        term.{
            | EString _ => term
            | EChar _ => term
            | EInt _ => term
            | EFloat _ => term
            | EVariable _ => term
            | EField e => e.EField(record = self.processTerm(functions, e.record))
            | EWildcard e => term
            | EList e => e.EList(
                items = e.items.map {| Pair(item, b) => Pair(self.processTerm(functions, item), b) }
            )
            | ESequential e => e.ESequential(
                before = self.processTerm(functions, e.before)
                after = self.processTerm(functions, e.after)
            )
            | ELet e =>
                let newFunctions = functions.remove(e.name)
                e.ELet(
                    value = self.processTerm(functions, e.value)
                    body = self.processTerm(newFunctions, e.body)
                )
            | ELambda e => e.ELambda(lambda = self.processLambda(functions, e.lambda))
            | EVariant e => e.EVariant(
                arguments = e.arguments.map { _.map { self.processArgument(functions, _) } }
            )
            | EVariantIs e => term
            | ECopy e => e.ECopy(
                record = self.processTerm(functions, e.record)
                arguments = e.arguments.map { self.processField(functions, _) }
            )
            | EPipe e => e.EPipe(
                value = self.processTerm(functions, e.value)
                function = self.processTerm(functions, e.function)
            )
            | ECall(at, StaticCall target, effect, typeArguments, arguments, _) {
                functions.get(target.name) | Some(signature)
            } =>
                let dictionaries = signature.constraints.map {
                    self.makeDictionary(at, signature.generics, typeArguments, _)
                }
                ECall(
                    at = at
                    target = target.StaticCall()
                    effect = effect
                    typeArguments = typeArguments
                    arguments = arguments.map { self.processArgument(functions, _) }
                    dictionaries = dictionaries
                )
            | ECall e =>
                let target = e.target.{
                    | DynamicCall call => call.DynamicCall(function = self.processTerm(functions, call.function))
                    | StaticCall _ => e.target
                }
                e.ECall(
                    target = target
                    arguments = e.arguments.map { self.processArgument(functions, _) }
                )
            | ERecord e => e.ERecord(fields = e.fields.map { self.processField(functions, _) })
            | EFunctions e =>
                let newFunctions = functions.addAll(e.functions.map { f =>
                    Pair(f.signature.name, f.signature)
                }.toMap())
                e.EFunctions(
                    functions = e.functions.map { self.processFunctionDefinition(newFunctions, _) }
                    body = self.processTerm(newFunctions, e.body)
                )
            | EAssign e => e.EAssign(value = self.processTerm(functions, e.value))
            | EAssignField e => e.EAssignField(
                record = self.processTerm(functions, e.record)
                value = self.processTerm(functions, e.value)
            )
        }
    }

    processArgument(functions: Map[String, Signature], argument: Argument): Argument {
        argument.Argument(value = self.processTerm(functions, argument.value))
    }

    processField(functions: Map[String, Signature], field: Field): Field {
        field.Field(value = self.processTerm(functions, field.value))
    }

    makeDictionary(
        at: Location,
        typeParameters: List[String],
        typeArguments: List[Type],
        constraint: Constraint
    ): Dictionary {
        let instantiationMap = typeParameters.zip(typeArguments).toMap()
        let unification = Unification.make([], False)
        let newGenerics = constraint.generics.map { unification.instantiate(instantiationMap, _) }
        let firstType = newGenerics.grabFirst().{
            | TConstructor t => t
            | TVariable t => fail(t.at, " is still a unification variable")
        }
        let instance = self.instances.get(InstanceKey(constraint.name, firstType.name)).else {
            throw(CompileError(at, "Missing instance " + firstType.name + ": " + constraint.name))
        }
        let dictionaries = instance.constraints.map { c =>
            self.makeDictionary(at, instance.generics, firstType.generics, c)
        }
        Dictionary(instance.packagePair, instance.moduleName, constraint.name, firstType.name, dictionaries)
    }

}

constraintsToInstances(constraints: List[Constraint]): Map[InstanceKey, InstanceValue] {
    constraints.map { c =>
        let typeName = c.generics.grabFirst().{
            | TConstructor(_, name, _) => name
            | TVariable(_, i) => fail(c.at, "Unexpected unification variable: $" + i)
        }
        Pair(
            InstanceKey(c.name, typeName)
            InstanceValue(
                generics = []
                constraints = []
                packagePair = PackagePair("", "")
                moduleName = ""
                traitName = c.name
                typeArguments = c.generics
            )
        )
    }.toMap()
}