sunday-wtf.hs

{-# Language AllowAmbiguousTypes #-}
{-# Language UndecidableInstances #-}
{-# Language TypeFamilyDependencies #-}
{-# Language MultiWayIf #-}
module Main where

import Data.Word
import Data.Bits
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Control.Monad.State
import Control.Monad.Trans.Maybe
import Data.ByteString.Builder
import Data.Maybe
import Lens.Micro.Platform
import Data.Kind
import GHC.TypeLits
import Data.Proxy
import Safe
import System.Random

data NOP
data LOADB
data SKIPBI
data ANDBI
data ORBI
data XORBI
data ADDBI
data SUBBI
data MULTBI
data REPEAT
data RET

class Emittable a where
  emit :: a -> Builder

class (Emittable (Arg a), KnownNat (Opcode a)) => Instruction a where
  type family Opcode a = (code :: Nat) | code -> a
  type family Arg a :: Type

data OP = forall a . (Instruction a, Emittable (Proxy a)) =>
  OP (Proxy a) (Arg a) | BYTE Word8

instance Instruction a => Emittable (Proxy a) where
  emit _ = word8 . fromIntegral $ natVal (Proxy @(Opcode a))

instance Emittable OP where
  emit (OP op arg) = emit op <> emit arg
  emit (BYTE w) = word8 w

instance Emittable () where
  emit = mempty

instance Emittable Word8 where
  emit = word8

instance Emittable b => Emittable [b]  where
  emit xs= mconcat (fmap emit xs)

instance Instruction NOP where
  type instance Opcode NOP = 0xFE
  type instance Arg NOP = ()

instance Instruction LOADB where
  type instance Opcode LOADB = 0x01
  type instance Arg LOADB = Word8

instance Instruction SKIPBI where
  type instance Opcode SKIPBI = 0x02
  type instance Arg SKIPBI = Word8

instance Instruction ORBI where
  type instance Opcode ORBI= 0x03
  type instance Arg ORBI = Word8

instance Instruction ANDBI where
  type instance Opcode ANDBI= 0x04
  type instance Arg ANDBI = Word8

instance Instruction XORBI where
  type instance Opcode XORBI= 0x05
  type instance Arg XORBI = Word8

instance Instruction ADDBI where
  type instance Opcode ADDBI = 0x06
  type instance Arg ADDBI = Word8

instance Instruction SUBBI where
  type instance Opcode SUBBI = 0x07
  type instance Arg SUBBI = Word8

instance Instruction MULTBI where
  type instance Opcode MULTBI = 0x08
  type instance Arg MULTBI = Word8

instance Instruction REPEAT where
  type instance Opcode REPEAT = 0xC0
  type instance Arg REPEAT = Word8

instance Instruction RET where
  type instance Opcode RET = 0xFF
  type instance Arg RET = ()

toOp :: Word8 -> OP
toOp = \case
  0x00 -> OP (Proxy @NOP) ()
  _    -> OP (Proxy @NOP) ()

op :: forall a . Instruction a
   => Arg a
   -> OP

op = OP (Proxy @a)

byte :: Word8 -> OP
byte = BYTE

wtf :: [ OP ]
wtf =  [ op @NOP ()
       , op @LOADB 66
       , op @SKIPBI 4
       , byte 0x01
       , byte 0x00
       , byte 0x80
       , byte 0x00
       , op @LOADB 5
       , op @ADDBI 61
       , op @SUBBI 66
       ]

runCode :: ByteString -> Maybe Word8
runCode s = execState (runMaybeT (go s)) Nothing
  where

    next = MaybeT . pure . BS.uncons

    go bs = do
      exit <- next bs >>= exec
      go exit

    exec (b, rest)
      | b == code @NOP    = nop rest
      | b == code @LOADB  = loadb rest
      | b == code @SKIPBI = skipbi rest
      | b == code @ORBI   = orbi rest
      | b == code @ANDBI  = andbi rest
      | b == code @XORBI  = xorbi rest
      | b == code @ADDBI  = addbi rest
      | b == code @SUBBI  = subbi rest
      | b == code @MULTBI = multi rest
      | b == code @REPEAT = repeatN rest
      | b == code @RET    = ret rest
      | otherwise         = nop rest

    ret _ = pure mempty

    nop = pure

    multi bs = do
      (n, rest) <- next bs
      apply (*) n
      pure rest

    addbi bs = do
      (n, rest) <- next bs
      apply (+) n
      pure rest

    subbi bs = do
      (n, rest) <- next bs
      apply (-) n
      pure rest

    orbi bs = do
      (n, rest) <- next bs
      apply (.|.) n
      pure rest

    andbi bs = do
      (n, rest) <- next bs
      apply (.&.) n
      pure rest

    xorbi bs = do
      (n, rest) <- next bs
      apply xor n
      pure rest

    skipbi bs = do
      (n, rest) <- next bs
      pure (BS.drop (fromIntegral n) rest)

    loadb bs = do
      (n, rest) <- next bs
      put (Just n)
      pure rest

    repeatN bs = do
      (n, rest) <- next bs

      rest' <- replicateM (min 16 (fromIntegral n)) $ do
                next rest >>= exec

      pure (lastDef "" rest')

    apply fn n = do
      st <- get
      put $ Just $ fromMaybe 0 st `fn` fromIntegral n

    code :: forall a b . (Integral b, Instruction a) => b
    code = fromIntegral (natVal (Proxy @(Opcode a)))

randomPrefix :: MonadIO m => m Builder
randomPrefix = liftIO do

  noise1 <- makeNoise

  rnum <- liftIO $ randomIO @Word8

  m1   <- liftIO $ randomRIO (0,2 :: Int)

  let rop = makeNum rnum m1

  m2 <- liftIO $ randomRIO (0,1 :: Int)

  let bs = toLazyByteString $ noise1 <> rop
                                     <> downTo0 rnum m2

  let rx = 16 - (LBS.length bs & fromIntegral)

  suff <- if rx <= 3 then
            replicateM rx (pure (op @NOP ())) <&> emit
          else
             makeNoiseN (fromIntegral rx - 2)

  pure $ lazyByteString bs <> suff

  where

    downTo0 rnum = \case
      0 -> emit [ op @ANDBI (complement rnum) ]
      1 -> emit [ op @XORBI rnum ]
      _ -> emit [ op @SUBBI rnum ]

    makeNum rnum = \case

      0 -> let (a,b) = rnum `divMod` 2 in emit [op @LOADB a, op @MULTBI 2, op @ADDBI b]

      1 | rnum < 4  -> emit [op @REPEAT rnum, op @ADDBI 1]

      _ -> emit $ op @LOADB rnum

    makeNoise = do
      randomRIO (0,4) >>= makeNoiseN

    makeNoiseN noiseN  = do

      if noiseN > 0 then do
        bytes <- replicateM (fromIntegral noiseN) (randomIO @Word8)
        pure $ emit ( op @SKIPBI noiseN : fmap byte bytes )
      else
        pure mempty


main :: IO ()
main = do
  print "okay"