Use current environment when evaluating functions

Fixes sharkdp#125.

src/Insect/Environment.purs

@@ -35,7 +35,7 @@ derive instance Eq StorageType
 data StoredValue = StoredValue StorageType Quantity
 
 -- | Mathematical functions on physical quantities.
-type MathFunction = NonEmpty List Quantity → Either EvalError Quantity
+type MathFunction = Environment → NonEmpty List Quantity → Either EvalError Quantity
 
 -- | Meta information for a function, mostly for pretty printing
 data FunctionDescription = BuiltinFunction (Maybe Int) | UserFunction (NonEmpty List Identifier) Expression
@@ -164,16 +164,16 @@ initialEnvironment =
     hiddenVal identifier value = Tuple identifier (StoredValue HiddenConstant value)
     constFunc  identifier func = Tuple identifier (StoredFunction Constant (wrapSimple identifier func) (BuiltinFunction (Just 1)))
     constFunc2 identifier func = Tuple identifier (StoredFunction Constant (wrapSimple2 identifier func) (BuiltinFunction (Just 2)))
-    constFuncN identifier func = Tuple identifier (StoredFunction Constant func (BuiltinFunction Nothing))
+    constFuncN identifier func = Tuple identifier (StoredFunction Constant (const func) (BuiltinFunction Nothing))
 
     wrapSimple ∷ Identifier → (Quantity → Either ConversionError Quantity) → MathFunction
-    wrapSimple name func qs =
+    wrapSimple name func _ qs =
       case qs of
         x :| Nil → lmap QConversionError $ func x
         _        → Left $ WrongArityError name 1 (length (NonEmptyList qs))
 
     wrapSimple2 ∷ Identifier → (Quantity → Quantity → Either ConversionError Quantity) → MathFunction
-    wrapSimple2 name func qs =
+    wrapSimple2 name func _ qs =
       case qs of
         x1 :| x2 : Nil → lmap QConversionError $ func x1 x2
         _              → Left $ WrongArityError name 2 (length (NonEmptyList qs))

src/Insect/Interpreter.purs

@@ -107,7 +107,7 @@ eval env (Apply name xs)        =
     else
       case lookup name env.functions of
         Just (StoredFunction _ fn _) →
-          traverse (eval env) xs >>= fn >>= checkFinite
+          traverse (eval env) xs >>= fn env >>= checkFinite
         Nothing → Left (LookupError name)
 eval env (BinOp op x y)         = do
   x' <- eval env x
@@ -286,7 +286,7 @@ runInsect env (FunctionAssignment name argNames expr) =
     numExpected = length argNames'
 
     userFunc ∷ MathFunction
-    userFunc argValues =
+    userFunc env' argValues =
       if numGiven == numExpected
         then evalAndSimplify functionEnv expr
         else Left (WrongArityError name numExpected numGiven)
@@ -296,8 +296,8 @@ runInsect env (FunctionAssignment name argNames expr) =
         args = zip argNames' argValues'
 
         insertArg map (Tuple argName val) = insert argName (StoredValue UserDefined val) map
-        functionEnv = env { values = foldl insertArg env.values args
-                          , functions = delete name env.functions
+        functionEnv = env' { values = foldl insertArg env'.values args
+                          , functions = delete name env'.functions
                           }
 
 runInsect env (PrettyPrintFunction name) =

test/Main.purs

@@ -889,6 +889,14 @@ main = runTest do
       expectOutput' "3 m" "3m"
       expectOutput' "3 m" " 3.0 meter  "
 
+    test "Function using current environment" do
+      let env1 = initialEnvironment
+          env2 = (repl fmtPlain env1 "a = 2").newEnv
+          env3 = (repl fmtPlain env2 "f(x) = a * x").newEnv
+          env4 = (repl fmtPlain env3 "a = 3").newEnv
+
+      expectOutput env4 "3" "f(1)"
+
     test "Square, cube and other exponentiation operators" do
       expectOutput' "18" "3²*2"
       expectOutput' "18" "3² 2"