{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE QuantifiedConstraints #-}
-- For Syn.LetRecable
{-# LANGUAGE UndecidableInstances #-}

module Symantic.Semantics.LetInserter where

import Control.Applicative (liftA2, liftA3)
import Control.Monad (Monad (..))
import Control.Monad.Trans.State.Strict qualified as MT
import Data.Bool (Bool (..), not, otherwise, (&&))
import Data.Either (Either (..))
import Data.Eq (Eq (..))
import Data.Foldable (foldMap)
import Data.Function (const, id, on, ($), (.))
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.Kind (Type)
import Data.List qualified as List
import Data.Maybe (Maybe (..), isNothing, maybe)
import Data.Ord (Ord (..))
import Data.Semigroup (Semigroup (..))
import Data.Set qualified as Set
import Debug.Trace qualified as Trace
import Numeric.Natural (Natural)
import Symantic.Syntaxes.Classes qualified as Syn
import Symantic.Syntaxes.Derive qualified as Syn
import Text.Show (Show (..), showListWith, showString, shows)
import Unsafe.Coerce (unsafeCoerce)
import Prelude (error, pred)
import Prelude qualified

-- import Control.Monad.Trans.Class qualified as MT
-- import Data.Functor.Identity (Identity(..))
-- import Data.Monoid (Monoid(..))
-- import Data.String (String)

traceShow :: Show a => a -> b -> b
-- traceShow = Trace.traceShow
traceShow _x a = a

newtype Locus = Locus Natural
  deriving (Eq, Ord)
  deriving newtype (Show)

{-
type Locus = [Pigit]
data Pigit = POne | PTwo | PThree | PLoc Locus | PNote String
  deriving (Eq, Ord)
instance Show Pigit where
  showsPrec _p x =
    case x of
      POne -> shows (1::Int)
      PTwo -> shows (2::Int)
      PThree -> shows (3::Int)
      PLoc loc -> shows loc
      PNote msg -> showString msg
-}

-- * Type 'OpenCode'
newtype OpenCode code a = OpenCode {unOpenCode :: Env code -> code a}
type instance Syn.Derived (OpenCode code) = code

-- instance Syn.Derivable (OpenCode code) where
--  derive oc = unOpenCode oc mempty
instance Syn.LiftDerived (OpenCode code) where
  liftDerived x = OpenCode (const x)
instance Syn.LiftDerived1 (OpenCode code) where
  liftDerived1 f x = OpenCode (\env -> f (unOpenCode x env))
instance Syn.LiftDerived2 (OpenCode code) where
  liftDerived2 f x y = OpenCode (\env -> f (unOpenCode x env) (unOpenCode y env))
instance Syn.LiftDerived3 (OpenCode code) where
  liftDerived3 f x y z = OpenCode (\env -> f (unOpenCode x env) (unOpenCode y env) (unOpenCode z env))

-- instance Syn.Abstractable code => Syn.Abstractable (OpenCode code) where
--  lam f = OpenCode (\env -> Syn.lam (\x -> unOpenCode f env (unOpenCode x env)))
-- instance Syn.Instantiable code => Syn.Instantiable (OpenCode code)
-- instance Syn.Constantable c code => Syn.Constantable c (OpenCode code)
instance Syn.IfThenElseable code => Syn.IfThenElseable (OpenCode code) where
  ifThenElse x y z = OpenCode (\env -> Syn.ifThenElse (unOpenCode x env) (unOpenCode y env) (unOpenCode z env))
instance Syn.Equalable code => Syn.Equalable (OpenCode code) where
  equal = OpenCode (const Syn.equal)
instance Syn.Constantable c code => Syn.Constantable c (OpenCode code)

-- ** Type Env'
type Env code = [EnvBind code]

-- *** Type 'Var'
newtype Var a = Var {unVar :: VarName}
  deriving (Eq)
  deriving newtype (Show)
type VarName = Locus

-- *** Type 'EnvBind'
data EnvBind code where
  EnvBind :: !(Var a, code a) -> EnvBind code
instance Show (EnvBind code) where
  showsPrec p (EnvBind (n, _e)) = showsPrec p n

lookupVar :: Var a -> OpenCode code a
lookupVar (Var v) =
  traceShow (["lookupVar", "OpenCode"], v) $
    OpenCode $ \env ->
      traceShow (["lookupVar", "OpenCode", "find"], v, ("env", env)) $
        case List.find (\(EnvBind (Var v1, _)) -> v1 == v) env of
          Just (EnvBind (_, x)) -> unsafeCoerce x
          Nothing -> error ("lookupVar failure for var=(" <> show v <> ")  env=" <> show ((\(EnvBind (Var n, _v)) -> n) <$> env))

lookupVarLetInserter :: Var a -> LetInserter code a
lookupVarLetInserter v =
  LetInserter $
    return
      AnnotatedCode
        { annotatedCodeOpenCode = lookupVar v
        , annotatedCodeFloatingLets = FloatingLets []
        , annotatedCodeFreeVars = Set.singleton (unVar v)
        }

insertVar :: Var a -> code a -> Env code -> Env code
insertVar v e env =
  traceShow ("insertVar", v) $
    EnvBind (v, e) : env

-- * Type 'Memo'
type Memo = Natural

-- ** Type 'MaybeMemo'

-- | Like a @(Maybe Memo)@ except 'Nothing' is never equal to any other value
-- to implement unique memoization key.
newtype MaybeMemo = MaybeMemo (Maybe Memo)
  deriving (Show)

instance Eq MaybeMemo where
  MaybeMemo (Just x) == MaybeMemo (Just y) = x == y
  _ == _ = False

-- * Type 'FloatingLet'
data FloatingLet code where
  FloatingLet ::
    { floatingLetLocus :: !(Maybe Locus)
    -- ^ 'FloatingLet' are either inserted at the place of an explicit 'Locus',
    -- or at the binding ('lam' or 'let_') that dominates all free variables of the bound expression,
    -- whichever has the narrowest scope.
    , floatingLetMemo :: !MaybeMemo
    -- ^ Memoization key, either internal (unique) or user-specified.
    --
    -- The memo key defines the equivalence classes:
    -- expressions with the same memo key are to be shared.
    -- Therefore, if @('genLetMemo' k l e)@ finds that there is already a let-binding
    -- produced by an earlier 'genLetMemo' with the same locus @(l)@ and the same memo key @(k)@,
    -- @('genLetMemo' k l e) returns the code of the earlier bound variable.
    , floatingLetVarName :: !(Var a)
    , floatingLetBoundCode :: !(BoundCode code a)
    , floatingLetFreeVars :: !(Set.Set VarName)
    } ->
    FloatingLet code
instance Show (FloatingLet code) where
  showsPrec _p FloatingLet{..} =
    showString "FloatingLet{"
      . showString "locus="
      . shows floatingLetLocus
      . showString ", "
      . showString "memo="
      . shows floatingLetMemo
      . showString ", "
      . showString "var="
      . shows floatingLetVarName
      . showString ", "
      . showString "exp="
      . showString (case floatingLetBoundCode of BoundCodeOpenCode{} -> "BoundCodeOpenCode"; BoundCodeLetInserter{} -> "BoundCodeLetInserter")
      . showString ", "
      . showString "fvs="
      . showListWith (shows) (Set.toList floatingLetFreeVars)
      . showString ", "
      . showString "}"

-- * Type 'VLBindings'

-- A virtual let-binding.
-- The key idea is virtual let-bindings: whereas clet introduces an ordinary let-binding
-- whose location is fixed, glet generates the code for a fresh variable accompanied by a virtual
-- binding of that variable. Virtual bindings do not have (yet) a fixed location: they are
-- attached to the expression that uses their bound variables and ‘float up’ when their attached
-- expression is incorporated into a bigger expression. Eventually, when they reach a point that
-- a metaprogrammer has marked with a dedicated primitive, virtual bindings are converted to
-- real let-bindings.
newtype FloatingLets code = FloatingLets {unFloatingLets :: [FloatingLet code]}
instance Show (FloatingLets code) where
  showsPrec _p (FloatingLets l) = go l
    where
      go [] = showString "\n"
      go (e : es) = showString "\n  , " . shows e . go es

-- | Append two lists skipping the duplicates.
-- The order of elements is preserved
mergeFloatingLets :: FloatingLets code -> FloatingLets code -> FloatingLets code
mergeFloatingLets (FloatingLets xs) (FloatingLets ys) =
  traceShow (["mergeFloatingLets"], xs, ys) $
    FloatingLets (loop (List.reverse xs) ys)
  where
    loop acc = \case
      [] -> List.reverse acc
      y@FloatingLet{floatingLetVarName = Var yv} : t ->
        case List.filter
          ( \a@FloatingLet{floatingLetVarName = Var av} ->
              floatingLetLocus y == floatingLetLocus a
                && floatingLetMemo y == floatingLetMemo a
                && yv == av
          )
          acc of
          [] -> loop (y : acc) t
          _ -> loop (acc) t

-- fs ->
--  loop (List.concatMap (\FloatingLet{floatingLetBoundCode, floatingLetLocus, floatingLetMemo, floatingLetVarName=Var av} ->
--    if av == yv
--    then
--    FloatingLet { floatingLetVarName=Var yv, .. }) fs <> acc) t
instance Semigroup (FloatingLets code) where
  (<>) = mergeFloatingLets

-- instance Monoid (FloatingLets code) where
--  mempty = FloatingLets []
--  mappend = (<>)

-- * Type 'BoundCode'

-- | Bound expression
data BoundCode (code :: Type -> Type) a
  = BoundCodeOpenCode !(OpenCode code a)
  | -- | In the recursive case,
    -- a virtual binding may carry a yet to be computed piece of code.
    BoundCodeLetInserter !(LetInserter code a)

-- * Type 'AnnotatedCode'

-- | A code value.
-- A code value is annotated only with the type of the expression it generates,
-- with no further classifiers (at present).
data AnnotatedCode (code :: Type -> Type) a = AnnotatedCode
  { annotatedCodeOpenCode :: !(OpenCode code a)
  , annotatedCodeFloatingLets :: !(FloatingLets code)
  , annotatedCodeFreeVars :: !(Set.Set VarName)
  -- ^ The free variables of the code,
  -- used to limit the floating up of locus-less 'FloatingLet's
  -- at 'lam' or 'let_' level, in order to avoid scope extrusion.
  }

-- annotatedCodeFreeVars :: AnnotatedCode code a -> Set.Set VarName
-- annotatedCodeFreeVars annotatedCode = foldMap floatingLetFreeVars (unFloatingLets (annotatedCodeFloatingLets annotatedCode))

runCD :: Syn.Letable code => AnnotatedCode code a -> code a
runCD AnnotatedCode{annotatedCodeOpenCode = oc, annotatedCodeFloatingLets = FloatingLets bindingFLS} =
  let bindingLFSByMemo = groupBy ((==) `on` floatingLetMemo) bindingFLS
  in (`unOpenCode` []) $ List.foldr bindFloatingLetMemo oc bindingLFSByMemo
runCD _ = error "runCD: virtual binding extrusion"

type instance Syn.Derived (AnnotatedCode code) = OpenCode code
instance Syn.LiftDerived (AnnotatedCode code) where
  liftDerived x = AnnotatedCode{annotatedCodeOpenCode = x, annotatedCodeFloatingLets = FloatingLets [], annotatedCodeFreeVars = Set.empty}
instance Syn.LiftDerived1 (AnnotatedCode code) where
  liftDerived1 f e1 =
    AnnotatedCode
      { annotatedCodeOpenCode = f (annotatedCodeOpenCode e1)
      , annotatedCodeFloatingLets = annotatedCodeFloatingLets e1
      , annotatedCodeFreeVars = annotatedCodeFreeVars e1
      }
instance Syn.LiftDerived2 (AnnotatedCode code) where
  liftDerived2 f e1 e2 =
    AnnotatedCode
      { annotatedCodeOpenCode = f (annotatedCodeOpenCode e1) (annotatedCodeOpenCode e2)
      , annotatedCodeFloatingLets = mergeFloatingLets (annotatedCodeFloatingLets e1) (annotatedCodeFloatingLets e2)
      , annotatedCodeFreeVars = annotatedCodeFreeVars e1 <> annotatedCodeFreeVars e2
      }
instance Syn.LiftDerived3 (AnnotatedCode code) where
  liftDerived3 f e1 e2 e3 =
    AnnotatedCode
      { annotatedCodeOpenCode = f (annotatedCodeOpenCode e1) (annotatedCodeOpenCode e2) (annotatedCodeOpenCode e3)
      , annotatedCodeFloatingLets = mergeFloatingLets (annotatedCodeFloatingLets e1) (mergeFloatingLets (annotatedCodeFloatingLets e2) (annotatedCodeFloatingLets e3))
      , annotatedCodeFreeVars = annotatedCodeFreeVars e1 <> annotatedCodeFreeVars e2 <> annotatedCodeFreeVars e3
      }
instance Syn.Constantable c code => Syn.Constantable c (AnnotatedCode code)
instance Syn.Instantiable code => Syn.Instantiable (AnnotatedCode code)
instance Syn.IfThenElseable code => Syn.IfThenElseable (AnnotatedCode code)
instance Syn.Equalable code => Syn.Equalable (AnnotatedCode code)

-- * Type 'LetInserter'
newtype LetInserter code a = LetInserter {unR :: MT.State Natural (AnnotatedCode code a)}
runLetInserter :: Syn.Letable code => LetInserter code a -> code a
runLetInserter = runCD . (`MT.evalState` 0) . unR

type instance Syn.Derived (LetInserter code) = AnnotatedCode code
instance Syn.LiftDerived (LetInserter sem) where
  liftDerived = LetInserter . return
instance Syn.LiftDerived1 (LetInserter sem) where
  liftDerived1 f e1 = LetInserter (f <$> unR e1)
instance Syn.LiftDerived2 (LetInserter sem) where
  liftDerived2 f e1 e2 = LetInserter (liftA2 f (unR e1) (unR e2))
instance Syn.LiftDerived3 (LetInserter sem) where
  liftDerived3 f e1 e2 e3 = LetInserter (liftA3 f (unR e1) (unR e2) (unR e3))
instance Syn.IfThenElseable code => Syn.IfThenElseable (LetInserter code)
instance Syn.Equalable code => Syn.Equalable (LetInserter code)
instance Syn.Constantable c code => Syn.Constantable c (LetInserter code)
instance Syn.Instantiable code => Syn.Instantiable (OpenCode code)

freshNatural :: MT.State Natural Natural
freshNatural = do
  n <- MT.get
  MT.put (Prelude.succ n)
  return n

-- type instance Syn.Derived (LetInserter sem) = Syn.Derived sem
freshVar :: MT.State Natural (Var a)
freshVar = Var . Locus <$> freshNatural

-- instance Syn.Constantable c sem => Syn.Constantable c (LetInserter sem) where
--   constant c = LetInserter $ pure $ Syn.constant c
instance Syn.Instantiable sem => Syn.Instantiable (LetInserter sem)

-- f .@ x = LetInserter $ Syn.liftDerived2 (Syn..@) (unR f) (unR x)
-- instance Syn.Abstractable sem => Syn.Abstractable (LetInserter sem) where
--   lam f = LetInserter $ \loc ->
--       Syn.lam $ \x ->
--       unR (f (LetInserter (\_loc -> x)))
--         (PNote "lam":loc)
-- instance Syn.Abstractable code => Syn.Abstractable (LetInserter (AnnotatedCode code)) where
--   lam body = LetInserter $ \loc ->
--     AnnotatedCode
--       { annotatedCodeOpenCode = _e
--       }
--     -- \x -> unR (body (LetInserter (\_ -> x))) (PNote "lam":loc)
-- instance Syn.Letable (LetInserter sem) where
--   let_ exp body = LetInserter $ \loc ->
--     let x = unR exp (POne:loc) in
--     unR (body (LetInserter (\_loc -> x))) (PTwo:loc)

-- instance Syn.Letable code => Syn.Letable (LetInserter code) where
--  let_ e body = LetInserter $ \loc ->
--    let x = unR e (POne:loc) in
--    unR (body (LetInserter (\_loc -> x))) (PTwo:loc)

-- instance Syn.Letable (LetInserter sem) where
--  let_ exp body = LetInserter (\loc -> let x = unR exp (POne:loc) in unR (body (LetInserter (\_loc -> x))) (PTwo:loc))
-- instance Syn.IfThenElseable code => Syn.IfThenElseable (LetInserter (AnnotatedCode code)) where
--  ifThenElse =
--    Syn.liftDerived3 $ \bC okC koC -> OpenCode $ \env ->
--      Syn.ifThenElse (unOpenCode bC env) (unOpenCode okC env) (unOpenCode koC env)
-- instance Syn.FromDerived Syn.Equalable sem => Syn.Equalable (LetInserter sem) where
--   equal = LetInserter (\loc -> trace ("equal "<>show loc) $ Syn.equal)

-- | Check to see there are no duplicate 'VarName's in 'VLBindings'.
-- This is just a safety check
check_no_dup_varname :: FloatingLets code -> a -> a
check_no_dup_varname = loop []
  where
    loop :: [VarName] -> FloatingLets code -> a -> a
    loop _found (FloatingLets []) = id
    loop found (FloatingLets (FloatingLet{floatingLetVarName = Var v, ..} : t))
      | v `List.elem` found =
          error ("Duplicate varname " <> show v <> ", with locus " <> show floatingLetLocus)
      | otherwise = loop (v : found) (FloatingLets t)

-- | let-bind all vlbindings with the given locus.
-- This is a non-recursive case. All these let-bindings are independent
-- It is up to the user to ensure the independence: if the bindings
-- are dependent, the user should have introduced several loci
-- Virtual let bindings are sort of free variables.
-- We bind free variables later. Likewise, we bind let-bindings later.
-- 'Memo' is ignored here.
-- class GenLetable sem where
-- instance Syn.Letable code => GenLetable (LetInserter (AnnotatedCode code)) where
floatingLetVarname :: FloatingLet code -> VarName
floatingLetVarname FloatingLet{floatingLetVarName = Var v} = v

instance
  ( Syn.Abstractable code
  , Syn.Letable code
  , Syn.LetRecable Int code
  ) =>
  Syn.Abstractable (LetInserter code)
  where
  lam body = LetInserter $ do
    vname <- freshVar
    unR $
      Syn.liftDerived1
        ( \bodyCD ->
            let (bindingFLS, passingFLS) =
                  List.partition
                    (\fl -> Set.member (unVar vname) (floatingLetFreeVars fl))
                    (unFloatingLets (annotatedCodeFloatingLets bodyCD))
            in let bindingLFSByMemo = groupBy ((==) `on` floatingLetMemo) bindingFLS
               in traceShow
                    ( "lam"
                    , ("vname", unVar vname)
                    , ("bindingFLS", FloatingLets bindingFLS)
                    , ("passingFLS", FloatingLets passingFLS)
                    , ("bodyFreeVars", annotatedCodeFreeVars bodyCD)
                    , ("annotatedCodeFloatingLets bodyCD", annotatedCodeFloatingLets bodyCD)
                    )
                    $ Syn.liftDerived1
                      ( \bodyOC ->
                          OpenCode $ \env ->
                            Syn.lam $ \x ->
                              (`unOpenCode` insertVar vname x env) $
                                List.foldr bindFloatingLetMemo bodyOC bindingLFSByMemo
                      )
                      bodyCD
                        { annotatedCodeFloatingLets = FloatingLets passingFLS
                        , annotatedCodeFreeVars =
                            Set.delete (unVar vname) $
                              annotatedCodeFreeVars bodyCD
                                `Set.difference` Set.fromList (floatingLetVarname <$> bindingFLS)
                        }
        )
        (body (lookupVarLetInserter vname))
instance
  ( Syn.Letable code
  , Syn.LetRecable Int code
  ) =>
  Syn.Letable (LetInserter code)
  where
  let_ expR body = LetInserter $ do
    vname <- freshVar
    unR $
      Syn.liftDerived2
        ( \expCD bodyCD ->
            let (bindingFLS, passingFLS) =
                  List.partition
                    (\fl -> Set.member (unVar vname) (floatingLetFreeVars fl))
                    (unFloatingLets (annotatedCodeFloatingLets bodyCD))
            in let bindingLFSByMemo = groupBy ((==) `on` floatingLetMemo) bindingFLS
               in Syn.liftDerived2
                    ( \expOC bodyOC ->
                        OpenCode $ \env ->
                          Syn.let_ (unOpenCode expOC env) $ \x ->
                            (`unOpenCode` (insertVar vname x env)) $
                              List.foldr bindFloatingLetMemo bodyOC bindingLFSByMemo
                    )
                    expCD
                    bodyCD
                      { annotatedCodeFloatingLets = FloatingLets passingFLS
                      , annotatedCodeFreeVars =
                          Set.delete (unVar vname) $
                            annotatedCodeFreeVars bodyCD
                              `Set.difference` Set.fromList (floatingLetVarname <$> bindingFLS)
                      }
        )
        expR
        (body (lookupVarLetInserter vname))

bindLet :: Syn.Letable code => FloatingLet code -> OpenCode code a -> OpenCode code a
bindLet fl2@FloatingLet{..} c =
  traceShow ("bindLet", fl2) $
    case floatingLetBoundCode of
      BoundCodeOpenCode oc -> OpenCode $ \env ->
        Syn.let_ (unOpenCode oc env) $ \x ->
          unOpenCode c $
            -- traceShow ("bindLet", floatingLetVarName) $
            insertVar floatingLetVarName x env
      _ -> error "bindLet is always called for evaluated bindings, without latent bound exp"

-- instance Syn.LetRecable idx (LetInserter sem) where
--  letRec _idx f body =
--    let self idx = LetInserter (\loc ->
--          --trace ("letRec "<>show loc) $
--          unR (f self idx) (PNote "fix":loc))
--    in body self

-- mletRec :: p -> ((t1 -> LetInserter sem a) -> t1 -> LetInserter sem a) -> ((t1 -> LetInserter sem a) -> t2) -> t2
-- mletRec _idx f body =
--   let self idx = LetInserter $
--         --trace ("letRec "<>show loc) $
--         unR (f self idx)
--   in body self
--
-- letR ::
--   Syn.LetRecable Int sem =>
--   (sem a -> sem a) ->
--   (sem a -> sem w) ->
--   sem w
-- letR f body = Syn.letRec (1::Int)
--   (\self _idx -> f (self 0))
--   (\self -> body (self 0))
-- type family MemoKey (sem:: Type -> Type) where MemoKey a = Memo

-- | Memoizing non-recursive genLet
class MemoGenLetable sem where
  withLocus :: (Locus -> sem a) -> sem a
  genLetMemoAtMaybe :: Maybe Memo -> Maybe Locus -> sem a -> sem a

genLetMemo :: MemoGenLetable sem => Memo -> Locus -> sem a -> sem a
genLetMemo m l = genLetMemoAtMaybe (Just m) (Just l)

genLet :: MemoGenLetable sem => sem a -> sem a
genLet = genLetMemoAtMaybe Nothing Nothing

genLetAt :: MemoGenLetable sem => Locus -> sem a -> sem a
genLetAt = genLetMemoAtMaybe Nothing . Just

instance Syn.Letable code => MemoGenLetable (LetInserter code) where
  withLocus body = LetInserter $ do
    lname <- Locus <$> freshNatural
    bodyCD <- unR (body lname)
    let fls = annotatedCodeFloatingLets bodyCD
    let bindingByLocusFLS = List.filter (\fl -> floatingLetLocus fl == Just lname) (unFloatingLets fls)
    let bindingVars = foldMap (Set.singleton . floatingLetVarname) bindingByLocusFLS
    let (bindingFLS, passingFLS) = List.partition (\fl -> not (Set.disjoint bindingVars (floatingLetFreeVars fl))) (unFloatingLets fls)
    let bindingLFSByMemo = groupBy ((==) `on` floatingLetMemo) bindingFLS
    return $
      traceShow (["withLocus"], ("bindingLFSByMemo", bindingLFSByMemo)) $
        -- check_no_dup_varname (annotatedCodeFloatingLets bodyCD) $
        AnnotatedCode
          { annotatedCodeOpenCode = List.foldr bindFloatingLetMemo (annotatedCodeOpenCode bodyCD) bindingLFSByMemo
          , annotatedCodeFloatingLets = FloatingLets passingFLS
          , annotatedCodeFreeVars = annotatedCodeFreeVars bodyCD `Set.difference` Set.fromList (floatingLetVarname <$> bindingFLS)
          }
  genLetMemoAtMaybe keyMaybe mlname body = LetInserter $ traceShow (["genLetMemoAtMaybe"]) $ do
    vname <- freshVar
    -- mkey <- runIdentity <$> unR key
    bodyCD <-
      traceShow ("genLetMemoAtMaybe", ("vname", vname), ("key", keyMaybe)) $
        unR body
    let fvs = Set.insert (unVar vname) (annotatedCodeFreeVars bodyCD)
    -- if the bodyression to let-bind has itself let-bindings,
    -- they are straightened out
    -- The new binding is added at the end; it may depend on
    -- the bindings in evl (as is the case for the nested genLet)
    return $
      AnnotatedCode
        { annotatedCodeOpenCode = lookupVar vname
        , annotatedCodeFloatingLets =
            FloatingLets $
              unFloatingLets (annotatedCodeFloatingLets bodyCD)
                <> [ FloatingLet
                      { floatingLetLocus = mlname
                      , floatingLetMemo = MaybeMemo keyMaybe
                      , floatingLetVarName = vname
                      , floatingLetBoundCode = BoundCodeOpenCode (annotatedCodeOpenCode bodyCD)
                      , floatingLetFreeVars = fvs
                      }
                   ]
        , annotatedCodeFreeVars = fvs
        }

-- Bind the variables in order, preserving the dependency.
-- Bind only one variable within each group, substituting the others
bindFloatingLetMemo :: Syn.Letable code => [FloatingLet code] -> OpenCode code a -> OpenCode code a
bindFloatingLetMemo gr oc =
  traceShow (["withLocus", "bindFloatingLetMemo"]) $
    case assertSameLocuses gr of
      [] -> error "bindFloatingLetMemo: empty group"
      [fl] -> bindLet fl oc
      fl@FloatingLet{floatingLetVarName = vName} : fls ->
        bindLet fl $ OpenCode $ \env ->
          let defCode = unOpenCode (lookupVar vName) env
          in let ext_equ env' fl2 =
                  reBind2 (memKey fl) defCode fl2
                    : env'
             in unOpenCode oc $ List.foldl' ext_equ env fls
  where
    memKey fl@FloatingLet{floatingLetBoundCode = BoundCodeOpenCode{}} = floatingLetMemo fl
    memKey fl = error $ "group pust be canonical: " <> show fl

assertSameLocuses :: [FloatingLet code] -> [FloatingLet code]
assertSameLocuses fls
  | Set.size (foldMap (Set.singleton . floatingLetLocus) fls) <= 1 = traceShow ("assertSameLocuses", fls) fls
  | otherwise = error "assertSameLocuses"

-- | Especially for 'Locus', all bindings have the same type.
-- In any event, 'Locus' plus the 'Memo' key witness for the bindings' type.
reBind2 :: MaybeMemo -> code a -> FloatingLet sem -> EnvBind code
reBind2 k c fl
  | k == floatingLetMemo fl =
      traceShow (["reBind2"], ("k", k), ("fl", fl)) $
        EnvBind (Var (floatingLetVarname fl), c)
reBind2 k _ fl = error $ "reBind2: locus/memo mismatch" <> show (("k", k), ("fl", fl))

{-
genLetAt :: LetInserter code a -> LetInserter code a
genLetAt = LetInserter $
  k <- freshKey
  genLetMemo
-}
-- TODO: Here genlet is a version of genlet l em e described earlier, for the case of all memo keys being distinct.
-- TODO: newtype LocusRec = Locus
{-
genLetM = LetInserter $ \loc ->
  k <- freshKey
  genLetMemo Nothing k
-}

-- | Unlike withLocus, we now have to evaluate expressions that
-- may be present in virtual bindings.
-- The rule: within an equivalence class, we evaluate only
-- the bexp of its representative (the first member of the group)
-- That evaluation may create more virtual bindings, which are to be merged

-- Pick vlbindings for the given locus and convert them to

-- * normal* equivalence classes.

-- Return the list of normal equivalence classes and the
-- remaining bindings (for different loci).
-- An equivalence class is called normal if its representative
-- (the first binding in the group) has the BoundCodeOpenCode as the bound
-- expression (that is, it is determined which future-stage expression to bind).
-- If a group is not normal, that is, its representative has a not
-- yet determined bound expression, we have to determine it --
-- which creates more bindings to normalize (some of which could be
-- members of existing groups -- but never replace the existing
-- representative -- we always preserve the order!)
--
-- Keeping the correct order and maintain dependency is actually
-- very difficult. When evaluating the expression bound to var v,
-- we obtain a code value with vlbindings. Some of them mention v
-- (and so should be put into the same group as v, but _after_ it),
-- some are new and so should be put in a group _before_ v because
-- the v binding may depend on them.
-- It seems simpler to just forget about the order between groups
-- and dependencies and generate one big letrec with mutually dependent clauses.

-- | @('groupBy' eq xs)@: partition a list @xs@
-- into a list of equivalence classes quotiented by the @eq@ function
--
-- The order of the groups follows the order in the input list;
-- within each group, the elements occur in the same order they
-- do in the input list.
groupBy :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy eq = go []
  where
    go acc [] = List.reverse acc
    go acc (x : xs) = go ((x : grp) : acc) rest
      where
        (grp, rest) = List.partition (eq x) xs

normalizeGroup :: Locus -> FloatingLets code -> MT.State Natural ([FloatingLets code], FloatingLets code)
normalizeGroup lname vl = do
  let (bindingFLS, passingFLS) = List.partition (\FloatingLet{..} -> floatingLetLocus == Just lname) (unFloatingLets vl)
  let bindingLFSByMemo = groupBy ((==) `on` floatingLetMemo) bindingFLS
  -- Try to find an abnormal group
  let (normalGroups, abnormalGroups) =
        let abnomality_check = \case
              FloatingLet{floatingLetBoundCode = BoundCodeLetInserter{}} : _ -> True
              _ -> False
        in List.break abnomality_check bindingLFSByMemo
  -- traceShow (["normalizeGroup"], ("lname",lname), ("bindingLFSByMemo", FloatingLets <$> bindingLFSByMemo), ("normalGroups", FloatingLets <$> normalGroups), ("abnormalGroups", FloatingLets <$> abnormalGroups), ("passingFLS", FloatingLets passingFLS)) $
  case abnormalGroups of
    -- all groups are normal
    [] ->
      return $
        traceShow (["normalizeGroup", "return"], ("lname", lname), ("normalGroups", FloatingLets <$> normalGroups), ("passingFLS", passingFLS)) $
          (FloatingLets <$> normalGroups, FloatingLets passingFLS)
    -- at least one group representative is abnormal
    (fl0@FloatingLet{floatingLetBoundCode = BoundCodeLetInserter fl0Exp, floatingLetVarName = fl0VarName} : fls0) : grps -> do
      fl0BodyCD <- unR fl0Exp
      -- The group has become normal
      let normalizedFL0 =
            FloatingLet
              { floatingLetLocus = floatingLetLocus fl0
              , floatingLetMemo = floatingLetMemo fl0
              , floatingLetVarName = fl0VarName
              , floatingLetBoundCode = BoundCodeOpenCode (annotatedCodeOpenCode fl0BodyCD)
              , -- Update the free variables, which could not be done in 'genLetMemoRec'
                -- where fl0BodyCD was not yet available.
                floatingLetFreeVars = floatingLetFreeVars fl0 <> annotatedCodeFreeVars fl0BodyCD
              }
              : fls0
      let groups = normalGroups <> (normalizedFL0 : grps)
      -- New bindings
      -- let vl' = FloatingLets (List.concat groups) <> FloatingLets passingFLS <> annotatedCodeFloatingLets fl0BodyCD
      let vl' = FloatingLets $ List.concat groups <> passingFLS <> unFloatingLets (annotatedCodeFloatingLets fl0BodyCD)
      normalizeGroup lname vl'
    _ -> error "normalizeGroup: groups are always non-empty"

-- * Class 'MemoGenLetRecable'
class MemoGenLetRecable sem where
  withLocusRec :: (Locus -> sem a) -> sem a
  genLetMemoRec :: Locus -> Memo -> sem a -> sem a

instance
  Syn.LetRecable Int code =>
  MemoGenLetRecable (LetInserter code)
  where
  withLocusRec body = LetInserter $ do
    lname <- Locus <$> freshNatural
    bodyCD <- unR $ body lname
    (bindingFLS, passingFLS) <-
      check_no_dup_varname (annotatedCodeFloatingLets bodyCD) $
        normalizeGroup lname (annotatedCodeFloatingLets bodyCD)
    -- Bind the variables in order, preserving the dependency
    return $
      traceShow (["withLocusRec", "return"], ("lname", lname), ("bindingFLS", bindingFLS), ("passingFLS", passingFLS)) $
        AnnotatedCode
          { annotatedCodeOpenCode = bindFloatingLetMemoRec bindingFLS (annotatedCodeOpenCode bodyCD)
          , annotatedCodeFloatingLets = passingFLS
          , annotatedCodeFreeVars = annotatedCodeFreeVars bodyCD `Set.difference` foldMap (Set.fromList . (floatingLetVarname <$>) . unFloatingLets) bindingFLS
          }
  genLetMemoRec lname key body = LetInserter $ do
    vname <- freshVar
    -- Cannot compute all the free variables yet
    -- because body is not yet evaluated to a 'AnnotatedCode'
    let fvs = Set.singleton (unVar vname)
    return $
      traceShow (["genLetMemoRec"], ("key", key), ("vname", vname)) $
        AnnotatedCode
          { annotatedCodeOpenCode = lookupVar vname
          , annotatedCodeFloatingLets =
              FloatingLets
                [ FloatingLet
                    { floatingLetLocus = Just lname
                    , floatingLetMemo = MaybeMemo (Just key)
                    , floatingLetVarName = vname
                    , floatingLetBoundCode = BoundCodeLetInserter body
                    , floatingLetFreeVars = fvs
                    }
                ]
          , annotatedCodeFreeVars = fvs
          }

-- | Create mutually-recursive letrec bindings.
-- All variables are bound at the same time.
-- The binding expression is taken from the group representative.
-- It has to be 'BoundCodeOpenCode': all groups assumed to be normal.
bindFloatingLetMemoRec ::
  Syn.LetRecable Int code =>
  [FloatingLets code] ->
  OpenCode code a ->
  OpenCode code a
bindFloatingLetMemoRec [] c = c
bindFloatingLetMemoRec bindingFLS c =
  traceShow (["bindFloatingLetMemoRec", "OpenCode"], ("bindingFLS", bindingFLS)) $
    OpenCode $ \env ->
      traceShow (["bindFloatingLetMemoRec", "Syn.letRec"], ("bindingFLS", bindingFLS), ("env", env)) $
        Syn.letRec
          len
          ( \self idx ->
              traceShow (["bindFloatingLetMemoRec", "Syn.letRec", "binds"], ("idx", idx)) $
                ( traceShow (["bindFloatingLetMemoRec", "Syn.letRec", "binds", "List.!!"], ("len", len), ("idx", idx), ("env", env)) $
                    unOpenCode (tobindOC List.!! idx)
                )
                  ( traceShow (["bindFloatingLetMemoRec", "Syn.letRec", "binds", "bind_all"], ("idx", idx), ("env", env)) $
                      bind_all self env
                  )
          )
          ( \self ->
              traceShow (["bindFloatingLetMemoRec", "Syn.letRec", "end"], ("env", env)) $
                unOpenCode c (bind_all self env)
          )
  where
    len = List.length bindingFLS
    -- (label,key) of group representative
    lk (fl@FloatingLet{floatingLetBoundCode = BoundCodeOpenCode{}} : _) = floatingLetMemo fl
    lk _ = error "groups must be non-empty and canonical"
    bind_group :: code a -> FloatingLets code -> [EnvBind code]
    -- Create bindings to bind all identifiers in a group
    bind_group ec (FloatingLets grp) =
      traceShow (["bindFloatingLetMemoRec", "bind_group"], ("grp", grp)) $
        List.map (reBind2 (lk grp) ec) (assertSameLocuses grp)
    bind_all self env =
      traceShow (["bindFloatingLetMemoRec", "bind_all"], ("bindingFLS", bindingFLS), ("env", env)) $
        List.concat $
          List.zipWith (\grp i -> bind_group (self i) grp) bindingFLS [0 :: Int .. pred len] <> [env]
    tobindOC =
      traceShow (["bindFloatingLetMemoRec", "tobindOC"], ("bindingFLS", bindingFLS)) $
        List.map
          ( \case
              -- TODO: avoid unwrapping
              FloatingLets (fl@FloatingLet{floatingLetBoundCode = BoundCodeOpenCode oc} : _) ->
                traceShow (["bindFloatingLetMemoRec", "tobindOC", "map"], fl) $
                  unsafeCoerce (oc)
              _ -> error "tobindOC"
          )
          bindingFLS

{-
genLetLocus :: Syn.Letable code => (Locus -> LetInserter code a) -> LetInserter code a
genLetLocus body = LetInserter $ do
  lname <- Locus <$> freshNatural
  bodyCD <-
    traceShow ("genLetLocus", ("lname", lname)) $
    unR (body lname)
  let (bindingFLS, passingFLS) =
        List.partition (\fl -> floatingLetLocus fl == Just lname)
                       (unFloatingLets (annotatedCodeFloatingLets bodyCD))
  -- Bind the variables in order, preserving the dependency
  --traceShow ("genLetLocus", (bindingFLS, passingFLS)) $
  return $
    check_no_dup_varname (annotatedCodeFloatingLets bodyCD) $
    AnnotatedCode{ annotatedCodeOpenCode = List.foldr bindLet (annotatedCodeOpenCode bodyCD) bindingFLS
      , annotatedCodeFloatingLets = FloatingLets passingFLS
      , annotatedCodeFreeVars = annotatedCodeFreeVars bodyCD `Set.difference` Set.fromList (floatingLetVarname <$> bindingFLS)
      }
-- In MetaOCaml, let(rec) is actually inserted either at the place
-- of the explicit let locus,
-- or at the binding that dominates all free variables of the bound expression,
-- whichever has the narrowest scope.
genLetAtMaybe :: Maybe Locus -> LetInserter code a -> LetInserter code a
genLetAtMaybe mlname body = LetInserter $ do
  vname <- freshVar
  bodyCD <-
        traceShow ("genLet", ("vname", vname)) $
        unR body
  -- if the expression to let-bind has itself let-bindings,
  -- they are straightened out
  -- The new binding is added at the end; it may depend on
  -- the bindings coming before (as is the case for the nested genLet)
  return $
    traceShow (["genLet", "bodyCD"], ("vname", vname), ("bodyCDFreeVars", (annotatedCodeFreeVars bodyCD))) $
    AnnotatedCode
      { annotatedCodeOpenCode = traceShow (["genLet", "lookupVar"], ("vname", vname)) $ lookupVar vname
      , annotatedCodeFloatingLets = FloatingLets $
          --(\vls -> traceShow (["genLet", "bindings"], vls) vls) $
          unFloatingLets (annotatedCodeFloatingLets bodyCD) <>
          [ FloatingLet
              { floatingLetLocus = mlname
              , floatingLetMemo = MaybeMemo (Just 0)
              , floatingLetVarName = vname
              , floatingLetBoundCode = BoundCodeOpenCode (annotatedCodeOpenCode bodyCD)
              , floatingLetFreeVars = Set.insert (unVar vname) (annotatedCodeFreeVars bodyCD)
              }
          ]
      , annotatedCodeFreeVars = Set.insert (unVar vname) $ annotatedCodeFreeVars bodyCD
      }
genLetAt :: Locus -> LetInserter code a -> LetInserter code a
genLetAt = genLetAtMaybe . Just

genLet :: LetInserter code a -> LetInserter code a
genLet = genLetAtMaybe Nothing
-}