open Mutton.Illuminate

open Mutton.Utils

open System.Collections.Generic

type Storage =

| S of string * int

member private this.DisplayText =

let (S (name, suffix)) = this

sprintf "%s.%d" name suffix

override this.ToString() = this.DisplayText

type Datum =

| DNum of int

| DSym of string

| DList of Datum list

type Bound =

| Var of Storage

| Def of Storage * Bound

| App of Bound * Bound list

| Fun of Storage * Bound

| Stx of Stx

| Quot of Datum

type BinderError = string * Firethorn.TextRange

type private BindingContext =

{

Parent: BindingContext option

Bindings: Dictionary<Ident, Storage>

}

/// Create an empty binding context

let empty = { Parent = None; Bindings = Dictionary<Ident, Storage>() }

/// Extend a binding context with a new name mapping

let extend ident bound ctx =

// Use indexer assignment to avoid exceptions on duplicate keys.

// This allows redefinition in the same scope without crashing.

ctx.Bindings.[ident] <- bound

/// Create a new child context with the given parent

let withParent parent =

{ Parent = Some parent; Bindings = Dictionary<Ident, Storage>() }

/// Resolve a name in the given context, searching parent contexts if

/// necessary. Returns the resolved name or the original name if not found.

let rec resolve ident ctx =

match ctx.Bindings.TryGetValue ident with

| true, bound -> bound

| false, _ ->

match ctx.Parent with

| Some parent -> resolve ident parent

| None -> S(ident.Name, 0)

let mutable private varSuffix = ref 0

let private freshStorage (name: string) : Storage =

let suffix = System.Threading.Interlocked.Increment(varSuffix)

S(name, suffix)

let rec private strip (stx: Stx) : Datum =

match stx with

| StxLiteral l ->

match l.Value with

| Some n -> DNum n

| None -> DSym "#err"

| StxIdent(id, _) -> DSym id.Name

| StxForm(items, _) -> DList(List.map strip items)

| StxClosure(stx, _) -> strip stx

let rec private bindOne (ctx: BindingContext) (stxEnv: StxEnv) =

function

| StxLiteral l ->

match l.Value with

| Some n -> Quot(DNum n) |> Some |> Ok

| None -> ($"Invalid literal %A{l}", l.Syntax.Range) |> Error

| StxIdent(id, _) ->

Var(BinderCtx.resolve id ctx) |> Some |> Ok

| StxForm(items, f) ->

match items with

| [] -> Error ($"No applicant in application form.", f.Syntax.Range)

| StxIdent(id, sym) :: args ->

match Expand.resolve id stxEnv with

| Illuminate.Lam -> bindLambda ctx stxEnv f args

| Illuminate.Def -> bindDefinition ctx stxEnv f args

| Illuminate.Quot -> bindQuotation f args

| Illuminate.Stx -> bindSyntaxQuotation f args

| _ -> bindSimpleApp ctx stxEnv f (StxIdent(id, sym)) args

| applicant :: args ->

bindSimpleApp ctx stxEnv f applicant args

| StxClosure(stx, env) -> bindOne ctx env stx

and private bindSimpleApp ctx stxEnv f applicant args =

bindOne ctx stxEnv applicant

|> Result.bind (fun called ->

args

|> List.map (bindOne ctx stxEnv)

|> accumulateResults

|> Result.bind (fun boundArgs ->

match called with

| Some called ->

App(called, (List.choose id boundArgs)) |> Some |> Ok

| None -> Error ($"Invalid function in application: has no value.", f.Syntax.Range)))

and private bindLambda ctx stxEnv f args =

match args with

| [ StxIdent(id, _); body ] ->

let argStorage = freshStorage id.Name

let bodyCtx =

BinderCtx.withParent ctx

BinderCtx.extend id argStorage bodyCtx

match bindOne bodyCtx stxEnv body with

| Ok(Some boundBody) ->

Fun(argStorage, boundBody) |> Some |> Ok

| Ok None ->

Error ($"Body has no value.", f.Syntax.Range)

| e -> e

| _ -> Error ($"Invalid lambda form %A{args}", f.Syntax.Range)

and private bindDefinition ctx stxEnv f args =

match args with

| [ StxIdent(id, _); body ] ->

match bindOne ctx stxEnv body with

| Ok(Some boundBody) ->

let newStore = freshStorage id.Name

BinderCtx.extend id newStore ctx

Def(newStore, boundBody) |> Some |> Result.Ok

| Ok None ->

Error ($"Invalid `def` body: has no value.", f.Syntax.Range)

| e -> e

| _ -> Error ($"Invalid `def` form %A{args}", f.Syntax.Range)

and private bindQuotation f args =

match args with

| [ form ] -> Quot(strip form) |> Some |> Result.Ok

| _ -> Error ($"Invalid `quot` form: %A{args}", f.Syntax.Range)

and private bindSyntaxQuotation f args =

match args with

| [ form ] -> Stx form |> Some |> Result.Ok

| _ -> Error ($"Invalid `stx` form: %A{args}", f.Syntax.Range)

let rec private bindSequence (ctx: BindingContext) (stxEnv: StxEnv) (exprs: Stx list) : Result<Bound list, string * Firethorn.TextRange> =

match exprs with

| [] -> Ok []

| stx :: rest ->

match Expand.expand stx stxEnv with

| None, newStxEnv ->

// def-syn: no runtime value produced; continue with updated env

bindSequence ctx newStxEnv rest

| Some expanded, newStxEnv ->

bindOne ctx newStxEnv expanded

|> Result.bind (fun bound ->

match bound with

| None -> bindSequence ctx newStxEnv rest

| Some b ->

bindSequence ctx newStxEnv rest

|> Result.map (fun bs -> b :: bs))

let public bind =

let ctx = BinderCtx.empty

bindSequence ctx Map.empty