{-# LANGUAGE ViewPatterns #-}
module Layout where

-- VLCurly is only inserted to get a column number of the first lexeme after
-- a layout starter since we don't keep full positions of every lexeme in Token
-- for clarity.
data Token =
    Token String
    | TokenVLCurly String !Int
    | TokenNL !Int
    deriving(Show)

data Context
    = NoLayout String String  -- what opened it and what we expect to close it.
    | Layout String !Int
    deriving(Show)

layout :: [Token] -> [Context] -> [Token]
layout (TokenNL n:Token "in":rs) (Layout "let" n':ls) = rbrace:Token "in":layout rs ls
layout (TokenNL n:Token s:rs) (Layout h n':ls)
    | s `elem` layoutContinuers = layout (Token s:rs) (Layout h (min n' n):ls)
layout (TokenNL n:rs) (Layout h n':ls)
    | n == n' = semi:layout rs (Layout h n':ls)
    | n > n' = layout rs (Layout h n':ls)
    | n < n' = rbrace:layout (TokenNL n:rs) ls
layout (TokenNL _:rs) ls = layout rs ls
layout (TokenVLCurly h n:rs) (Layout h' n':ls)
    | n > n' = lbrace:layout rs (Layout h n:Layout h' n':ls)
    | otherwise = lbrace : rbrace : layout rs (Layout h' n':ls)
layout (TokenVLCurly h n:rs) ls = lbrace:layout rs (Layout h n:ls)
layout (t@(Token s):rs) (dropLayouts -> (n,Just (b,e),ls)) | s == e
    = replicate n rbrace ++ t:layout rs ls
layout (t@(Token s):rs) ls | Just e <- lookup s layoutBrackets
    = t:layout rs (NoLayout s e:ls)
layout (t@(Token s):rs) ls@(Layout c _:_) |
    Just e <- lookup c conditionalBrackets >>= lookup s = t:layout rs (NoLayout s e:ls)
layout (t@(Token "in"):rs) ls = case ls of
    Layout "let" n:ls -> rbrace:t:layout rs ls
    ls -> t:layout rs ls
layout (t@(Token ","):rs) (Layout "let" _:NoLayout "|" e:ls) = rbrace:layout (t:rs) (NoLayout "|" e:ls)
layout (t@(Token "where"):rs) ls = case ls of
    Layout l n : rest | l `elem` ["do","of"]
        -> rbrace : t : layout rs rest -- 'where' closes 'do' and 'case' on equal indentation.
    _otherwise -> t : layout rs ls

layout (t:rs) ls = t:layout rs ls
layout [] (Layout _ n:ls) = rbrace:layout [] ls
layout [] [] = []
layout x y = error $ "unexpected layout: " ++ show (x,y)

-- unwind all pending layouts
dropLayouts :: [Context] -> (Int,Maybe (String,String),[Context])
dropLayouts cs = f 0 cs where
    f n [] = (n,Nothing,[])
    f n (NoLayout b e:ls) = (n,Just (b,e),ls)
    f n (Layout {}:ls) = f (n + 1) ls

semi = Token ";"
lbrace = Token "{"
rbrace = Token "}"

fsts = map fst
snds = map snd

layoutStarters   = ["where","let","of","do"]

-- these symbols will never close a layout.
layoutContinuers = ["|","->","=",";",","]

-- valid in all contexts
layoutBrackets = [
    ("case","of"),
    ("if","then"),
    ("then","else"),
    ("(",")"),
    ("[","]"),
    ("{","}")
    ]

conditionalBrackets = [
    ("of",[("|","->")]),
    ("let",[("|","=")]),
    ("[",[("|","]")])
    ]