{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

-- | Extended instructions are the representation I wish to use for the optimization techniques in these modules.
-- They should be easy to extend, hence I factored the 'Operation' type out of 'Language.Brainfuck.CompressedInstruction.CompressedInstruction'.

module Language.Brainfuck.Instruction.Extended (ExtendedInstruction(Modify, Move, Interact, Jump), pattern IfNonZero, pattern WithOffset, mkIfNonZero, mkWithOffset, translationSize, parse) where

import Data.Word (Word8)

import Data.Vector (Vector)
import qualified Data.Vector as Vector

import Numeric.Natural (Natural)

import Language.Brainfuck.Instruction.Compressed (CompressedInstruction)
import Language.Brainfuck.Instruction.Extended.Operation (Operation)
import Language.Brainfuck.Instruction.Extended.Interaction (Interaction)


import qualified Language.Brainfuck.Instruction.Compressed as CompressedInstruction
import qualified Language.Brainfuck.Instruction.Extended.Operation as Operation
import qualified Language.Brainfuck.Instruction.Extended.Interaction as Interaction

-- | This pattern protects the invariant that 'WithOffset'-Instructions may not be nested
--
-- >>> WithOffset 5 (Interact Interaction.Write)
-- AtOffset 5 (Interact Write)
--
-- >>> WithOffset 5 (WithOffset (-2) (Modify Operation.Add 5))
-- AtOffset 3 (Modify Add 5)

pattern WithOffset :: Integer -> ExtendedInstruction -> ExtendedInstruction
pattern WithOffset offset embedded <- AtOffset offset embedded
  where
    WithOffset offset embedded = mkWithOffset offset embedded

-- | This pattern protects the invariant that 'IfNonZero'-Instructions may not be directly nested
--
-- >>> IfNonZero (Modify Operation.Add 5)
-- WhenNonZero (Modify Add 5)
--
-- >>> IfNonZero (IfNonZero (Modify Operation.Subtract 2))
-- WhenNonZero (Modify Subtract 2)

pattern IfNonZero :: ExtendedInstruction -> ExtendedInstruction
pattern IfNonZero instruction <- WhenNonZero instruction
  where
    IfNonZero instruction = mkIfNonZero instruction

-- | The extended instruction set is defined by this enum.
--
-- Note: not all instruction types are exported since some of them hold invariants. 
--
-- You need to additionally consider these pattern exports when matching:
--
--   [@WithOffset@]: Will merge nested WithOffset instructions
--
--   [@IfNonZero@]: Will eliminate doubly conditional instructions

data ExtendedInstruction
  = AtOffset Integer ExtendedInstruction -- invariant, WithOffset may not nest itself
  | WhenNonZero ExtendedInstruction -- invariant, IfNonZero may not nest itself
  | Modify Operation Word8
  | Move Integer
  | Interact Interaction
  | Jump Integer
  deriving (Show)

-- | Smart constructor alias for the pattern 'WithOffset'
-- 
-- >>> mkIfNonZero $ mkIfNonZero (Interact Interaction.Read)
-- WhenNonZero (Interact Read)

mkIfNonZero :: ExtendedInstruction -> ExtendedInstruction
mkIfNonZero = \case
  WhenNonZero i -> WhenNonZero i
  instruction -> WhenNonZero instruction

-- | Smart constructor alias for the pattern 'IfNonZero'
--
-- >>> mkWithOffset 15 $ mkWithOffset (-3) $ Interact Interaction.Read
-- AtOffset 12 (Interact Read)

mkWithOffset :: Integer -> ExtendedInstruction -> ExtendedInstruction
mkWithOffset offset = \case
  AtOffset offset' i -> AtOffset (offset + offset') i
  instruction -> AtOffset offset instruction

-- | Parsess the CompressedInstruction program into the ExtendedInstruction form, which e.g. does not contain 'CompressedInstruction.Loop' instructions anymore.
--
-- >>> :set -XOverloadedLists
-- >>> parse [CompressedInstruction.Add 5]
-- [Modify Add 5]
--
-- >>> parse  [CompressedInstruction.Add 5, CompressedInstruction.Loop [CompressedInstruction.Loop [CompressedInstruction.ReadByte], CompressedInstruction.PutByte]]
-- [Modify Add 5,Jump 4,Jump 1,Interact Read,WhenNonZero (Jump (-2)),Interact Write,WhenNonZero (Jump (-5))]

parse :: Vector CompressedInstruction -> Vector ExtendedInstruction
parse = Vector.fromList . Vector.foldr prependTranslation []

-- | This prepends the translation of the instruction to the given list, this was before I learned about List.concatMap being very efficient

prependTranslation :: CompressedInstruction -> [ExtendedInstruction] -> [ExtendedInstruction]
prependTranslation instruction rest = let
  addSingle = (:rest)
  in case instruction of
    CompressedInstruction.Add       i    -> addSingle $ Modify Operation.Add i
    CompressedInstruction.Subtract  i    -> addSingle $ Modify Operation.Subtract i
    CompressedInstruction.MoveRight n    -> addSingle $ Move $ toInteger n
    CompressedInstruction.MoveLeft  n    -> addSingle $ Move $ negate . toInteger $ n
    CompressedInstruction.ReadByte       -> addSingle $ Interact Interaction.Read
    CompressedInstruction.PutByte        -> addSingle $ Interact Interaction.Write
    CompressedInstruction.Loop      body -> let
      bodySize = translationSize body
      backJump = IfNonZero $ Jump $ -(toInteger bodySize + 1)
      forwardJump = Jump $ toInteger bodySize

      in forwardJump : Vector.foldr prependTranslation (backJump : rest) body

-- | Calculates the required size for a instruction vector of ExtendedInstruction, since they may need more space.

translationSize :: Vector CompressedInstruction -> Natural
translationSize = Vector.sum . Vector.map instructionTranslationSize
  where
    instructionTranslationSize = \case
     CompressedInstruction.Add       _    -> 1
     CompressedInstruction.Subtract  _    -> 1
     CompressedInstruction.MoveRight _    -> 1
     CompressedInstruction.MoveLeft  _    -> 1
     CompressedInstruction.ReadByte       -> 1
     CompressedInstruction.PutByte        -> 1
     CompressedInstruction.Loop      body -> let
      skipLoopSize = 1
      repeatLoopSize = 1
      bodySize = translationSize body
      in skipLoopSize + repeatLoopSize + bodySize