-- TODO: These case statements could be combined
infer :: Expr -> Infer (Type, [Constraint])
infer expr =
  case expr of
    Lit (LInt _) -> return (typeInt, [])
    Lit (LBool _) -> return (typeBool, [])
    Var x -> do
      t <- lookupEnv x
      return (t, [])
    Lam x e -> do
      tv <- fresh
      (t, c) <- inEnv (x, Forall [] tv) (infer e)
      return (tv `TArr` t, c)
    App e1 e2 -> do
      (t1, c1) <- infer e1
      (t2, c2) <- infer e2
      tv <- fresh
      return (tv, c1 ++ c2 ++ [(t1, t2 `TArr` tv)])
    Let x e1 e2 -> do
      (t1, c1) <- infer e1
      env <- ask
      case runSolve c1 of
        Left err -> throwError err
        Right sub -> do
          let sc = generalize (apply sub env) (apply sub t1)
          (t2, c2) <- inEnv (x, sc) $ local (apply sub) (infer e2)
          return (t2, c1 ++ c2)
    Fix e -> do
      (t, c) <- infer e
      tv <- fresh
      return (tv, c ++ [(tv `TArr` tv, t)])
    Op op lhs rhs -> do
      (t1, c1) <- infer lhs
      (t2, c2) <- infer rhs
      tv <- fresh
      let u1 = t1 `TArr` (t2 `TArr` tv)
          u2 = ops op
      return (tv, c1 ++ c2 ++ [(u1, u2)])
    If cond tr fl -> do
      (t1, c1) <- infer cond
      (t2, c2) <- infer tr
      (t3, c3) <- infer fl
      return (t2, c1 ++ c2 ++ c3 ++ [(t1, typeBool), (t2, t3)])
infer :: Expr -> Infer (Type, [Constraint])
infer expr =
  case expr of
    Lit (LInt _) -> return (typeInt, [])
    Lit (LBool _) -> return (typeBool, [])
    Var x -> do
      t <- lookupEnv x
      return (t, [])
    Lam x e -> do
      tv <- fresh
      (t, c) <- inEnv (x, Forall [] tv) (infer e)
      return (tv `TArr` t, c)
    App e1 e2 -> do
      (t1, c1) <- infer e1
      (t2, c2) <- infer e2
      tv <- fresh
      return (tv, c1 ++ c2 ++ [(t1, t2 `TArr` tv)])
    Let x e1 e2 -> do
      (t1, c1) <- infer e1
      env <- ask
      case runSolve c1 of
        Left err -> throwError err
        Right sub -> do
          let sc = generalize (apply sub env) (apply sub t1)
          (t2, c2) <- inEnv (x, sc) $ local (apply sub) (infer e2)
          return (t2, c1 ++ c2)
    Fix e -> do
      (t, c) <- infer e
      tv <- fresh
      return (tv, c ++ [(tv `TArr` tv, t)])
    Op op lhs rhs -> do
      (t1, c1) <- infer lhs
      (t2, c2) <- infer rhs
      tv <- fresh
      let u1 = t1 `TArr` (t2 `TArr` tv)
          u2 = ops op
      return (tv, c1 ++ c2 ++ [(u1, u2)])
    If cond tr fl -> do
      (t1, c1) <- infer cond
      (t2, c2) <- infer tr
      (t3, c3) <- infer fl
      return (t2, c1 ++ c2 ++ c3 ++ [(t1, typeBool), (t2, t3)])
infer expr =
  case expr of
    Lit (LInt _) -> return (typeInt, [])
    Lit (LBool _) -> return (typeBool, [])
    Var x -> do
      t <- lookupEnv x
      return (t, [])
    Lam x e -> do
      tv <- fresh
      (t, c) <- inEnv (x, Forall [] tv) (infer e)
      return (tv `TArr` t, c)
    App e1 e2 -> do
      (t1, c1) <- infer e1
      (t2, c2) <- infer e2
      tv <- fresh
      return (tv, c1 ++ c2 ++ [(t1, t2 `TArr` tv)])
    Let x e1 e2 -> do
      (t1, c1) <- infer e1
      env <- ask
      case runSolve c1 of
        Left err -> throwError err
        Right sub -> do
          let sc = generalize (apply sub env) (apply sub t1)
          (t2, c2) <- inEnv (x, sc) $ local (apply sub) (infer e2)
          return (t2, c1 ++ c2)
    Fix e -> do
      (t, c) <- infer e
      tv <- fresh
      return (tv, c ++ [(tv `TArr` tv, t)])
    Op op lhs rhs -> do
      (t1, c1) <- infer lhs
      (t2, c2) <- infer rhs
      tv <- fresh
      let u1 = t1 `TArr` (t2 `TArr` tv)
          u2 = ops op
      return (tv, c1 ++ c2 ++ [(u1, u2)])
    If cond tr fl -> do
      (t1, c1) <- infer cond
      (t2, c2) <- infer tr
      (t3, c3) <- infer fl
      return (t2, c1 ++ c2 ++ c3 ++ [(t1, typeBool), (t2, t3)])