Simple.hs

{-# LANGUAGE OverloadedStrings #-}

module Zwirn.Language.Simple
  ( simplify,
    simplifyLoc,
    SimpleTerm (..),
    SimpleDef (..),
    LocSimpleTerm,
  )
where

{-
    Simple.hs - desugaring of the zwirn AST
    Copyright (C) 2023, Martin Gius

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this library.  If not, see <http://www.gnu.org/licenses/>.
-}

import Data.Text as Text (Text, filter, pack)
import Zwirn.Language.Location
import Zwirn.Language.Syntax

type LocSimpleTerm = Located SimpleTerm

-- simple representation of patterns
data SimpleTerm
  = SVar !Var
  | SText !Text
  | SNum !Text
  | SRest
  | SSeq [LocSimpleTerm]
  | SStack [LocSimpleTerm]
  | SChoice Int [LocSimpleTerm]
  | SLambda Var LocSimpleTerm
  | SApp LocSimpleTerm LocSimpleTerm
  | SInfix LocSimpleTerm LocVar LocSimpleTerm
  | SBracket LocSimpleTerm
  deriving (Eq, Show)

data SimpleDef
  = LetS !Var LocSimpleTerm
  deriving (Eq, Show)

simplifyLoc :: LocTerm -> LocSimpleTerm
simplifyLoc (Located p t) = Located p (simplify t)

simplify :: Term -> SimpleTerm
simplify (TVar x) = SVar x
simplify (TText x) = SText $ stripText x
  where
    stripText = Text.filter (/= '\"')
simplify (TNum x) = SNum x
simplify TRest = SRest
simplify x@(TRepeat _ _) = SSeq $ map (fmap simplify) $ resolveRepeat (Located NoLoc x)
simplify (TSeq ts) = SSeq (map (fmap simplify) $ concatMap resolveRepeat ts)
simplify (TStack ts) = SStack (map (fmap simplify) ts)
simplify (TChoice i ts) = SChoice i (map (fmap simplify) ts)
simplify (TAlt ts) = SBracket $ noLoc $ SInfix (noLoc $ SSeq ss) (noLoc "(/)") (noLoc $ SNum (pack $ show $ length ss))
  where
    ss = map (fmap simplify) $ concatMap resolveRepeat ts
simplify (TPoly (Located _ (TSeq ts)) n) = SBracket $ noLoc $ SInfix (noLoc $ SInfix (noLoc $ SSeq ss) (noLoc "(/)") (noLoc $ SNum (pack $ show $ length ss))) (noLoc "(*)") (fmap simplify n)
  where
    ss = map (fmap simplify) $ concatMap resolveRepeat ts
simplify (TPoly x n) = SInfix (fmap simplify x) (noLoc "(*)") (fmap simplify n)
simplify (TLambda [] t) = simplify (lValue t)
simplify (TLambda (x : xs) t) = SLambda x (noLoc $ simplify $ TLambda xs t)
simplify (TIfThenElse x y (Just z)) = SApp (noLoc $ SApp (noLoc $ SApp (noLoc $ SVar "ifthen") (fmap simplify x)) (fmap simplify y)) (fmap simplify z)
simplify (TIfThenElse x y Nothing) = SApp (noLoc $ SApp (noLoc $ SVar "if") (fmap simplify x)) (fmap simplify y)
simplify (TApp x y) = SApp (fmap simplify x) (fmap simplify y)
simplify (TInfix x op y) = SInfix (fmap simplify x) op (fmap simplify y)
simplify (TSectionR op y) = SLambda "_x" (noLoc $ SInfix (noLoc $ SVar "_x") op (fmap simplify y))
simplify (TSectionL x op) = SLambda "_x" (noLoc $ SInfix (fmap simplify x) op (noLoc $ SVar "_x"))
simplify (TBracket x) = SBracket (fmap simplify x)
simplify (TEnum Run x y) = SApp (noLoc $ SApp (noLoc $ SVar "runFromTo") (fmap simplify x)) (fmap simplify y)
simplify (TEnumThen Run x y z) = SApp (noLoc $ SApp (noLoc $ SApp (noLoc $ SVar "runFromThenTo") (fmap simplify x)) (fmap simplify y)) (fmap simplify z)
simplify (TEnum Alt x y) = SApp (noLoc $ SApp (noLoc $ SVar "slowrunFromTo") (fmap simplify x)) (fmap simplify y)
simplify (TEnumThen Alt x y z) = SApp (noLoc $ SApp (noLoc $ SApp (noLoc $ SVar "slowrunFromThenTo") (fmap simplify x)) (fmap simplify y)) (fmap simplify z)
simplify (TEnum Cord x y) = SApp (noLoc $ SApp (noLoc $ SVar "cordFromTo") (fmap simplify x)) (fmap simplify y)
simplify (TEnumThen Cord x y z) = SApp (noLoc $ SApp (noLoc $ SApp (noLoc $ SVar "cordFromThenTo") (fmap simplify x)) (fmap simplify y)) (fmap simplify z)
simplify (TEnum Choice x y) = SApp (noLoc $ SApp (noLoc $ SVar "chooseFromTo") (fmap simplify x)) (fmap simplify y)
simplify (TEnumThen Choice x y z) = SApp (noLoc $ SApp (noLoc $ SApp (noLoc $ SVar "chooseFromThenTo") (fmap simplify x)) (fmap simplify y)) (fmap simplify z)
simplify _ = error "Found unsubstituted macro while desugaring!"

resolveRepeat :: LocTerm -> [LocTerm]
resolveRepeat t = case getTotalRepeat t of
  Located _ (TRepeat x (Just i)) -> replicate i x
  Located _ (TRepeat x Nothing) -> [x, x]
  x -> [x]

-- TODO : not completely right when Nothing followed by Just...
getRepeat :: (LocTerm, Int) -> LocTerm
getRepeat (Located _ (TRepeat x (Just j)), k) = getRepeat (x, j * k)
getRepeat (Located _ (TRepeat x Nothing), k) = getRepeat (x, k + 1)
getRepeat (x, j) = noLoc $ TRepeat x (Just j)

getTotalRepeat :: LocTerm -> LocTerm
getTotalRepeat (Located _ (TRepeat t (Just i))) = getRepeat (t, i)
getTotalRepeat (Located _ (TRepeat t Nothing)) = getRepeat (t, 2)
getTotalRepeat t = t