-- MIT License, Copyright (c) 2024 Marvin Borner
-- Based on Benn Lynn's ION machine and John Tromp's nf.c
-- TODO: clean everything up after tests are working
--   iter -> fold/etc, application infix vm abstraction,
--   monadic VM, map -> array?
-- fairly literate translation from C code; TODO: use more functional style
{-# LANGUAGE NamedFieldPuns #-}
module Reducer.ION
  ( reduce
  ) where

import           Data.Bits                      ( (.|.) )
import           Data.Char                      ( chr
                                                , ord
                                                )
import qualified Data.Map                      as M
import           Data.Map                       ( Map )

import           Helper

ncomb :: Int
ncomb = 128

spTop :: Int
spTop = maxBound

data VM = VM
  { sp   :: Int
  , hp   :: Int
  , nvar :: Int
  , mem  :: Map Int Int
  }
  deriving Show

isComb :: Int -> Bool
isComb n = n < ncomb

-- default chr panics at high n
chr' :: Int -> Char
chr' n | n < 256 = chr n
chr' _           = '?'

new :: VM
new = VM spTop ncomb 0 mempty

load :: Int -> VM -> Int
load k vm = M.findWithDefault 0 k (mem vm)

store :: Int -> Int -> VM -> VM
store k v vm = vm { mem = M.insert k v $ mem vm }

app :: Int -> Int -> VM -> (Int, VM)
app x y vm@VM { hp } = (hp, store hp x $ store (hp + 1) y $ vm { hp = hp + 2 })

arg' :: Int -> VM -> Int
arg' n vm = load (load (sp vm + n) vm + 1) vm

arg :: Int -> VM -> (Int, VM)
arg n vm = (arg' n vm, vm)

apparg :: Int -> Int -> VM -> (Int, VM)
apparg m n vm = app (arg' m vm) (arg' n vm) vm

wor :: a -> b -> (a, b)
wor = (,)

com :: Char -> b -> (Int, b)
com = wor . ord

lazy :: Int -> (VM -> (Int, VM)) -> (VM -> (Int, VM)) -> VM -> VM
lazy d f a vm = do
  let
    fix i _ | i > d = Nothing
    fix i vm' =
      if load (load (sp vm' + i + 1) vm') vm' == load (sp vm' + i) vm'
        then fix (i + 1) vm'
        else do
          let cnvar = nvar vm'
          let
            (n1, vm1') =
              app (ord 'V') cnvar (vm' { sp = sp vm' - 2 }) { nvar = cnvar + 1 }
          let spi2       = load (sp vm1' + i + 2) vm1'
          let (n2, vm2') = app spi2 n1 (store (sp vm1' + i) spi2 vm1')
          let (n3, vm3') = app (ord 'L') n2 (store (sp vm2' + i + 1) n2 vm2')
          let vm4'       = store (sp vm3' + i + 2) n3 vm3'
          let vm5'       = store (load (sp vm4' + i + 3) vm4' + 1) n3 vm4'
          Just $ vm5' { sp = sp vm5' + i }
  case fix 1 vm of
    Just vm' -> vm'
    Nothing  -> do
      let (f', vm1) = f vm
      let (a', vm2) = a vm1
      let vm3       = vm2 { sp = sp vm + d }
      let dst       = load (sp vm + d + 1) vm
      store (sp vm3) f' (store dst f' (store (dst + 1) a' vm3))

rules :: Int -> VM -> VM
rules ch vm = case chr' ch of
  'M' -> lvm 0 (arg 1) (arg 1)
  'Y' -> lvm 0 (arg 1) (app (ord 'Y') 1)
  'I' -> if arg' 2 vm == load (sp vm + 1) vm
    then lvm 1 (arg 1) (arg 1)
    else lvm 1 (arg 1) (arg 2)
  'F' -> do
    let xv = arg' 2 vm
    if isComb xv
      then lvm 1 (com 'I') (arg 2)
      else lvm 1 (wor $ load xv vm) (wor $ load (xv + 1) vm)
  'f' -> do
    let x = load (load (sp vm + 2) vm + 1) vm
    let (v, vm') = if isComb x
          then do
            let cnvar = nvar vm
            let (a, vm1) =
                  app (ord 'V') cnvar vm { sp = sp vm + 1, nvar = cnvar + 1 }
            let (b, vm2) = app x a vm1
            let (c, vm3) = app (ord 'L') b vm2
            (c, store (load (sp vm3 + 1) vm3 + 1) c vm3)
          else (x, vm { sp = sp vm + 1 })
    store (sp vm') v vm'
  'K' -> do
    let (xv, _) = arg 1 vm
    if isComb xv
      then lvm 1 (com 'I') (wor xv)
      else lvm 1 (wor $ load xv vm) (wor $ load (xv + 1) vm)
  'T' -> lvm 1 (arg 2) (arg 1)
  'D' -> lvm 2 (arg 1) (arg 2)
  'B' -> lvm 2 (arg 1) (apparg 2 3)
  'C' -> lvm 2 (apparg 1 3) (arg 2)
  'R' -> lvm 2 (apparg 2 3) (arg 1)
  ':' -> lvm 2 (apparg 3 1) (arg 2)
  'S' -> lvm 2 (apparg 1 3) (apparg 2 3)
  'L' -> store (sp vm) (arg' 1 vm) vm
  'V' -> do
    let iter vm'@VM { sp } | sp == spTop = vm' { sp = -1 }
        iter vm'@VM { sp } | load (load (sp + 1) vm') vm' == load sp vm' = vm'
        iter vm'@VM { sp } =
          let parent = load (sp + 1) vm'
              left   = load parent vm'
              vm''   = if left == ord 'L'
                then vm' { nvar = nvar vm' - 1 }
                else store parent (load (left + 1) vm') vm'
          in  iter vm'' { sp = sp + 1 }
    let vm1 = iter vm { sp = sp vm + 1 }
    if sp vm1 == -1
      then vm1
      else do
        let parentVal = load (load (sp vm1 + 1) vm1 + 1) vm1
        let (a, vm2)  = app (ord 'f') (load (sp vm1) vm1) vm1
        lazy 0 (wor a) (wor parentVal) (store (sp vm2) parentVal vm2)
  '\0' -> store (sp vm) (arg' 1 vm) vm
  _    -> error $ "invalid combinator " ++ show ch
  where lvm n f g = lazy n f g vm

eval :: VM -> VM
eval vm@VM { sp } | sp == -1 = vm { sp = spTop }
eval vm@VM { sp } =
  let x1  = load sp vm
      x2  = load x1 vm
      vm1 = store (sp - 1) x2 vm { sp = sp - 1 }
      x3  = load x2 vm1
      vm2 = store (sp - 2) x3 vm1 { sp = sp - 2 }
  in  if isComb x1
        then eval $ rules x1 vm
        else if isComb x2
          then eval $ rules x2 vm1
          else if isComb x3 then eval $ rules x3 vm2 else eval vm2

run :: Int -> VM -> VM
run i vm@VM { sp } = eval $ store sp i vm

hasVar0 :: Int -> Int -> VM -> Bool
hasVar0 db depth vm = do
  let f = load db vm
  let a = load (db + 1) vm
  case chr' f of
    'V' -> a == depth
    'L' -> hasVar0 a (depth + 1) vm
    _   -> hasVar0 f depth vm || hasVar0 a depth vm

eta :: Int -> VM -> (Int, VM)
eta x vm = do
  let f = load x vm
  let a = load (x + 1) vm
  if not (isComb f) && load a vm == ord 'V' && load (a + 1) vm == 0 && not
       (hasVar0 f 0 vm)
    then (f, vm)
    else app (ord 'L') x vm

dbIndex :: Int -> Int -> VM -> (Int, VM)
dbIndex x depth vm = do
  let f = load x vm
  let a = load (x + 1) vm
  case chr' f of
    'V' -> app f (depth - 1 - a) vm
    'L' -> do
      let (idx, vm1) = dbIndex a (depth + 1) vm
      eta idx vm1
    _ -> do
      let (f', vm1) = dbIndex f depth vm
      let (a', vm2) = dbIndex a depth vm1
      app f' a' vm2

clapp :: (Int, Int) -> VM -> (Int, VM)
clapp (f, a) vm = case (chr' f, chr' a) of
  ('K', 'I')                      -> vord 'F'
  ('B', 'K')                      -> vord 'D'
  ('C', 'I')                      -> vord 'T'
  ('D', 'I')                      -> vord 'K'
  (_, 'I') | load f vm == ord 'B' -> (load (f + 1) vm, vm)
  (_, 'I') | load f vm == ord 'R' -> app (ord 'T') (load (f + 1) vm) vm
  (_, 'T') | load f vm == ord 'B' && load (f + 1) vm == ord 'C' -> vord ':'
  (_, 'I') | load f vm == ord 'S' && load (f + 1) vm == ord 'I' -> vord 'M'
  _                               -> app f a vm
  where vord c = (ord c, vm)

unDoubleVar :: Int -> Int -> VM -> (Int, VM)
unDoubleVar n db vm = do
  let f  = load db vm
  let a  = load (db + 1) vm
  let qf = load f vm == ord 'V' && load (f + 1) vm == n
  let qa = load a vm == ord 'V' && load (a + 1) vm == n
  if f == ord 'V'
    then (db, vm)
    else if f == ord 'L'
      then
        let (uda, vm') = unDoubleVar (n + 1) a vm
        in  if uda == 0 then (0, vm') else app (ord 'L') uda vm'
      else if qf && qa
        then app (ord 'V') n vm
        else if qf || qa
          then (0, vm)
          else
            let (udf, vm' ) = unDoubleVar n f vm
                (uda, vm'') = unDoubleVar n a vm'
            in  if udf == 0
                  then (0, vm')
                  else if uda == 0 then (0, vm'') else app udf uda vm''

recursive :: Int -> Int -> VM -> (Int, VM)
recursive f a vm = do
  let f'         = load (a + 1) vm
  let (f'', vm') = unDoubleVar 0 f' vm
  -- logical subtleties!
  if f == ord 'M' && load a vm == ord 'L' && load f' vm /= ord 'V'
    then if f'' /= 0 then app (ord 'L') f'' vm' else (0, vm')
    else (0, vm)

combineK :: (Int, Int, Int, Int) -> VM -> (Int, VM)
combineK (n1, d1, n2, d2) vm = do
  if n1 == 0
    then if n2 == 0
      then clapp (d1, d2) vm
      else if load (n2 + 1) vm /= 0
        then
          let (n, vm1) = clapp (ord 'B', d1) vm
          in  combineK (0, n, load n2 vm1, d2) vm1
        else combineK (0, d1, load n2 vm, d2) vm
    else if load (n1 + 1) vm /= 0
      then if n2 == 0
        then
          let (n, vm1) = clapp (ord 'R', d2) vm
          in  combineK (0, n, load n1 vm1, d1) vm1
        else if load (n2 + 1) vm /= 0
          then
            let (n, vm1) = combineK (0, ord 'S', load n1 vm, d1) vm
            in  combineK (load n1 vm, n, load n2 vm1, d2) vm1
          else
            let (n, vm1) = combineK (0, ord 'C', load n1 vm, d1) vm
            in  combineK (load n1 vm, n, load n2 vm1, d2) vm1
      else if n2 == 0
        then combineK (load n1 vm, d1, 0, d2) vm
        else if load (n2 + 1) vm /= 0
          then if load n2 vm == 0 && load (n2 + 1) vm /= 0 && d2 == ord 'I'
            then (d1, vm) -- eta optimization
            else
              let (n, vm1) = combineK (0, ord 'B', load n1 vm, d1) vm
              in  combineK (load n1 vm, n, load n2 vm1, d2) vm1
          else combineK (load n1 vm, d1, load n2 vm, d2) vm

zipN :: Int -> Int -> VM -> (Int, VM)
zipN nf na vm | nf == 0 = (na, vm)
zipN nf na vm | na == 0 = (nf, vm)
zipN nf na vm           = do
  let (z, vm1) = zipN (load nf vm) (load na vm) vm
  app z ((load (nf + 1) vm1) .|. (load (na + 1) vm1)) vm1

convertK :: Int -> VM -> (Int, Int, VM)
convertK db vm = do
  let f = load db vm
  let a = load (db + 1) vm
  case chr' f of
    'V' -> do
      let iter n 0 vm' = (n, vm')
          iter n i vm' = let (n', vm1) = app n 0 vm' in iter n' (i - 1) vm1
      let (nf, vm1) = let (n, vm') = app 0 1 vm in iter n a vm'
      (nf, ord 'I', vm1)
    'L' -> do
      let (na, ca, vm1) = convertK a vm
      let pn            = load na vm1
      if na == 0
        then let (n, vm2) = clapp (ord 'K', ca) vm1 in (0, n, vm2)
        else if load (na + 1) vm1 == 0
          then let (n, vm2) = combineK (0, ord 'K', pn, ca) vm1 in (pn, n, vm2)
          else (pn, ca, vm1)
    _ -> do
      let (nf, cf, vm1) = convertK f vm
      let (cf', a', vm2) =
            let (ca', vm') = recursive cf a vm1
            in  if nf == 0
                  then if ca' /= 0 then (ord 'Y', ca', vm') else (cf, a, vm')
                  else (cf, a, vm1)
      let (na, ca, vm3) = convertK a' vm2
      let (pn, vm4)     = zipN nf na vm3
      let (n, vm5)      = combineK (nf, cf', na, ca) vm4
      (pn, n, vm5)

toCLK :: Int -> VM -> (Int, VM)
toCLK db vm = do
  let (n, cl, vm1) = convertK db vm
  if n == 0 then (cl, vm1) else error "term not closed"

resolveExpression :: Int -> VM -> Expression
resolveExpression n _ | isComb n = error "unexpected combinator"
resolveExpression n vm           = do
  let f = load n vm
  let a = load (n + 1) vm
  case chr' f of
    'V' -> Bruijn a
    'L' -> Abstraction $ resolveExpression a vm
    _   -> Application (resolveExpression f vm) (resolveExpression a vm)

parseExpression :: Expression -> VM -> (Int, VM)
parseExpression (Abstraction t) vm = do
  let (t', vm1) = parseExpression t vm
  app (ord 'L') t' vm1
parseExpression (Application l r) vm = do
  let (l', vm1) = parseExpression l vm
  let (r', vm2) = parseExpression r vm1
  app l' r' vm2
parseExpression (Bruijn i) vm = do
  app (ord 'V') i vm
parseExpression _ _ = error "invalid expression"

reduce :: Expression -> Expression
reduce e = do
  let vm         = new
  let (db, vm1) = parseExpression (Abstraction e) vm
  let (cl, vm2)  = toCLK db vm1
  let res        = run cl vm2
  let (idx, fin) = dbIndex (load spTop res) 0 res
  case resolveExpression idx fin of
    Abstraction t -> t
    t             -> error $ "unexpected result " ++ show t