content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE StandaloneKindSignatures #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- | Unsafe coercions for linearly typed code.
  --
  -- Use this module to coerce non-linear functions to be linear or values
  -- bound linearly to be another type. /All/ functions in this module are
  -- unsafe.
  --
  -- Hence:
  --
  -- * Import this module qualifed as Unsafe.
  -- * Do not use this unless you have to. Specifically, if you can write a
  -- linear function @f :: A %1-> B@, do not write a non-linear version and coerce
  -- it.
  module Unsafe.Linear
    ( -- * Unsafe Coercions
      coerce,
      toLinear,
      toLinear2,
      toLinear3,
      toLinearN,
      ToLinearN (..),
    )
  where
  
  import Data.Kind (Constraint)
  import Data.Type.Equality (type (~~))
  import GHC.Exts (RuntimeRep (..), TYPE)
  import GHC.TypeNats
  import Unsafe.Coerce (UnsafeEquality (..), unsafeEqualityProof)
  
  -- | Linearly typed @unsafeCoerce@
  coerce :: forall a b. a %1 -> b
  coerce a =
    case unsafeEqualityProof @a @b of
      UnsafeRefl -> a
  {-# INLINE coerce #-}
  
  -- | Converts an unrestricted function into a linear function
  toLinear ::
    forall
      (r1 :: RuntimeRep)
      (r2 :: RuntimeRep)
      (a :: TYPE r1)
      (b :: TYPE r2)
      p
      x.
    (a %p -> b) %1 ->
    (a %x -> b)
  toLinear f = case unsafeEqualityProof @p @x of
    UnsafeRefl -> f
  
  -- | Like 'toLinear' but for two-argument functions
  toLinear2 ::
    forall
      (r1 :: RuntimeRep)
      (r2 :: RuntimeRep)
      (r3 :: RuntimeRep)
      (a :: TYPE r1)
      (b :: TYPE r2)
      (c :: TYPE r3)
      p
      q
      x
      y.
    (a %p -> b %q -> c) %1 ->
    (a %x -> b %y -> c)
  toLinear2 f = case unsafeEqualityProof @'(p, q) @'(x, y) of
    UnsafeRefl -> f
  
  -- | Like 'toLinear' but for three-argument functions
  toLinear3 ::
    forall
      (r1 :: RuntimeRep)
      (r2 :: RuntimeRep)
      (r3 :: RuntimeRep)
      (r4 :: RuntimeRep)
      (a :: TYPE r1)
      (b :: TYPE r2)
      (c :: TYPE r3)
      (d :: TYPE r4)
      p
      q
      r
      x
      y
      z.
    (a %p -> b %q -> c %r -> d) %1 ->
    (a %x -> b %y -> c %z -> d)
  toLinear3 f = case unsafeEqualityProof @'(p, q, r) @'(x, y, z) of
    UnsafeRefl -> f
  
  -- | @toLinearN@ subsumes the functionality of 'toLinear1', 'toLinear2', and
  -- 'toLinear3'. In particular, @toLinearN \@n@ unsafely changes the
  -- multiplicities of the first @n@ arrows from any multiplicity to any
  -- other multiplicity. To be explicit about how each multiplicity is
  -- being changed, you can use additional type arguments.
  --
  -- === Examples
  --
  -- @
  -- toLinearN \@2 :: (a %m-> b %n-> Int) %1-> a %x-> b %y-> Int
  -- toLinearN \@3 \@(_ %m-> _ -> _ %1-> _) \@(_ %1-> _ %1-> _ %x-> _)
  --   :: (a %m-> b -> c %1-> d) %1-> (a %1-> b %1-> c %x-> d)
  -- 'toLinear3' = toLinearN \@3
  -- @
  toLinearN :: forall n f g. (ToLinearN n f g) => f %1 -> g
  -- See Note: Core size
  toLinearN f = case unsafeLinearityProofN @n @f @g of
    UnsafeRefl -> f
  
  -- | @ToLinearN n f g@ means that @f@ and @g@ are the same with the
  -- possible exception of the multiplicities of the first @n@ arrows.
  type ToLinearN :: forall {rep :: RuntimeRep}. Nat -> TYPE rep -> TYPE rep -> Constraint
  class ToLinearN n f g where
    -- | Given that @f@ and @g@ are the same, with the possible exception of the
    -- multiplicities of the first @n@ arrows, @unsafeLinearityProofN \@n \@f \@g@
    -- is a fake proof that @f@ and @g@ are identical. This is used primarily in the
    -- definition of 'toLinearN', but it can also be used, for example, to coerce
    -- a container of functions:
    --
    -- @
    -- linearMany :: forall a b c. [a -> b -> c] %1-> [a %1-> b %1-> c]
    -- linearMany = castWithUnsafe (applyUnsafe (UnsafeRefl @[]) $
    --   unsafeLinearityProofN @2 @(a -> b -> c) @(a %1-> b %1-> c))
    --
    -- applyUnsafe :: UnsafeEquality f g -> UnsafeEquality x y -> UnsafeEquality (f x) (g y)
    -- applyUnsafe UnsafeRefl UnsafeRefl = UnsafeRefl
    --
    -- castWithUnsafe :: UnsafeEquality x y -> x %1-> y
    -- castWithUnsafe UnsafeRefl x = x
    -- @
    --
    -- The rather explicit handling of coercions seems to be necessary,
    -- unfortunately, presumably due to the way GHC eagerly rejects equality
    -- constraints it sees as definitely unsatisfiable.
    unsafeLinearityProofN :: UnsafeEquality f g
  
  instance (ToLinearN' ni f g, ni ~ ToINat n) => ToLinearN n f g where
    unsafeLinearityProofN = prf @ni
  
  -- | Plain old inductive natural numbers.
  data INat = Z | S INat
  
  -- | Convert a GHC 'Nat' to a real inductive natural number.
  -- We use this because GHC 'Nat' offers a friendly API but
  -- it's a terrible pain for implementation.
  type ToINat :: Nat -> INat
  type family ToINat n where
    ToINat 0 = 'Z
    ToINat n = 'S (ToINat (n - 1))
  
  -- | The actual implementation of 'ToLinearN', using the inductive natural
  -- number it's handed.
  type ToLinearN' :: forall {rep :: RuntimeRep}. INat -> TYPE rep -> TYPE rep -> Constraint
  class ToLinearN' arrs f g where
    prf :: UnsafeEquality f g
  
  instance (a ~ b) => ToLinearN' 'Z (a :: TYPE rep) (b :: TYPE rep) where
    prf = UnsafeRefl
  
  -- We use heterogeneous equality here to shift @rep ~ 'LiftedRep@ to the left
  -- side of the fat arrow, so that seeing @'S k@ lets GHC /infer/ that @x@ and
  -- @y@ are lifted types, rather than needing that information to come from
  -- elsewhere.
  instance
    ( ToLinearN' k fb gb,
      x ~~ ((a :: TYPE repa) %p -> (fb :: TYPE repb)),
      y ~~ (a %q -> (gb :: TYPE repb))
    ) =>
    ToLinearN' ('S k) (x :: TYPE rep) (y :: TYPE rep)
    where
    prf = case prf @k @fb @gb of
      UnsafeRefl -> case unsafeEqualityProof @p @q of
        UnsafeRefl -> UnsafeRefl
  
  -- Note: Core size
  --
  -- In GHC 9.0, at least, using toLinearN to implement toLinear2 and toLinear3
  -- produces more Core terms, though fewer types and coercions. Does this hamper
  -- inlining? If not, we could actually use it so. Alternatively, we *might* be
  -- able to reduce unsafeLinearityProofN to one case expression by using some
  -- (quantified constraint?) tricks. But that will be complicated, if possible,
  -- and probably not worth the trouble.
  