diff --git a/arkoala-arkts/libarkts/plugins/input/memo.sts b/arkoala-arkts/libarkts/plugins/input/memo.sts
index c655abe05732f0cab9208363fdc1f74f758fff6c..d1390606726e591bc93e9568255f7f0bdd93bc9e 100644
--- a/arkoala-arkts/libarkts/plugins/input/memo.sts
+++ b/arkoala-arkts/libarkts/plugins/input/memo.sts
@@ -20,47 +20,41 @@ class Test {
 	 	return arg
 	}
 
-	// TODO: context.compute<T>
-	// @memo method_with_type_parameter<T>(arg: T): T {
-	// 	return arg
-	// }
+	@memo method_with_type_parameter<T>(arg: T): T {
+		return arg
+	}
 
-	// TODO: return statement duplicates
-	// @memo static static_method_with_type_parameter<T>(arg: T): void {
-	// 	return
-	// }
+    @memo static static_method_with_type_parameter<T>(arg: T): void {
+		return
+	}
 
 	@memo internal_call() {
-		// TODO: the call doesn't rewrite
-		// this.void_method()
+		this.void_method()
 	}
 
-	//@memo lambda_arg(@memo arg: () => void) {
+	@memo lambda_arg(@memo arg: () => void) {
 
-	//}
+	}
 
-	//@memo lambda_arg_with_arg(@memo arg: (value: string) => string) {
+	@memo lambda_arg_with_arg(@memo arg: (value: string) => string) {
 
-	//}
+	}
 }
 
 class Use {
 	test() {
 		const test = new Test()
 
-		// TODO: the calls don't rewrite
-		// test.void_method()
-		// test.void_method_with_arg("an arg")
-		// test.void_method_with_return("a value")
-		// Test.static_method_with_type_parameter("I'm static")
+		test.void_method()
+		test.void_method_with_arg("an arg")
+		test.void_method_with_return("a value")
+		Test.static_method_with_type_parameter("I'm static")
 
-		// TODO: the function don't rewrite
-		// test.string_method_with_return("a string")
-		// test.method_with_type_parameter("I'm string")
+		test.string_method_with_return("a string")
+		test.method_with_type_parameter("I'm string")
 
-		// TODO: lamda calls
-		// test.lambda_arg(() => {})
-		// test.lambda_arg((value: string) => value)
+		test.lambda_arg(() => {})
+		test.lambda_arg_with_arg((value: string) => value)
 	}
 }
 
diff --git a/arkoala-arkts/memo-plugin/demo/demo.sts b/arkoala-arkts/memo-plugin/demo/demo.sts
index ff83381120d557bbbc46f6d3d8e81a7ae704bf7c..8e49383b95c5b31af7ff9109c978ce932711b6bc 100644
--- a/arkoala-arkts/memo-plugin/demo/demo.sts
+++ b/arkoala-arkts/memo-plugin/demo/demo.sts
@@ -3,7 +3,6 @@ import { memo_foo } from "./stub"
 
 @memo
 function f(s: string) {
-    console.log(s)
     let x = () => {
         console.log(s)
     }
diff --git a/arkoala-arkts/memo-plugin/src/FunctionTransformer.ts b/arkoala-arkts/memo-plugin/src/FunctionTransformer.ts
index 47e8ad344710d0aa851ab75f3cbcae6d32381e69..93217012915ae13f6616ee418d07f42186c422c3 100644
--- a/arkoala-arkts/memo-plugin/src/FunctionTransformer.ts
+++ b/arkoala-arkts/memo-plugin/src/FunctionTransformer.ts
@@ -23,8 +23,8 @@ import {
 import { ParameterTransformer } from "./ParameterTransformer"
 import { ReturnTransformer } from "./ReturnTranformer"
 
-function hasMemoAnnotation(node: arkts.ScriptFunction) {
-    return arkts.getAnnotations(node).some((it) =>
+function hasMemoAnnotation(node: arkts.ScriptFunction | arkts.ETSParameterExpression) {
+    return node.annotations.some((it) =>
         it.expr !== undefined && arkts.isIdentifier(it.expr) && it.expr.name === RuntimeNames.ANNOTATION
     )
 }
@@ -140,11 +140,17 @@ export class FunctionTransformer extends AbstractVisitor {
             const expr = node.expression
             const decl = arkts.getDecl(expr)
             if (decl instanceof arkts.MethodDefinition && hasMemoAnnotation(decl.scriptFunction)) {
+                const updatedArguments = node.arguments.map((it, index) => {
+                    if (hasMemoAnnotation(decl.scriptFunction.parameters[index])) {
+                        return factory.createComputeExpression(this.positionalIdTracker.id(decl.name.name), it)
+                    }
+                    return it
+                })
                 return arkts.factory.updateCallExpression(
                     node,
                     node.expression,
                     undefined,
-                    [...factory.createHiddenArguments(this.positionalIdTracker.id(decl.name.name)), ...node.arguments]
+                    [...factory.createHiddenArguments(this.positionalIdTracker.id(decl.name.name)), ...updatedArguments]
                 )
             }
         }
diff --git a/arkoala-arkts/memo-plugin/src/MemoFactory.ts b/arkoala-arkts/memo-plugin/src/MemoFactory.ts
index e4760ac18d28eb529142c1e5f40c0f2a18098c94..5c9246d91f3f493cf0a58997d1ec827b81324c0a 100644
--- a/arkoala-arkts/memo-plugin/src/MemoFactory.ts
+++ b/arkoala-arkts/memo-plugin/src/MemoFactory.ts
@@ -172,4 +172,35 @@ export class factory {
             syntheticReturnStatement,
         )
     }
+
+    // Compute
+    static createLambdaWrapper(node: arkts.AstNode): arkts.ArrowFunctionExpression {
+        return arkts.factory.createArrowFunction(
+            arkts.factory.createScriptFunction(
+                arkts.factory.createBlock([
+                    arkts.factory.createReturnStatement(node)
+                ]),
+                arkts.Es2pandaScriptFunctionFlags.SCRIPT_FUNCTION_FLAGS_ARROW,
+                arkts.Es2pandaModifierFlags.MODIFIER_FLAGS_NONE,
+                false,
+                undefined,
+                [],
+                undefined,
+                undefined,
+            )
+        )
+    }
+    static createComputeExpression(hash: arkts.NumberLiteral | arkts.StringLiteral, node: arkts.AstNode): arkts.CallExpression {
+        return arkts.factory.createCallExpression(
+            arkts.factory.createMemberExpression(
+                arkts.factory.createIdentifier(RuntimeNames.CONTEXT),
+                arkts.factory.createIdentifier(RuntimeNames.COMPUTE),
+                arkts.Es2pandaMemberExpressionKind.MEMBER_EXPRESSION_KIND_PROPERTY_ACCESS,
+                false,
+                false
+            ),
+            undefined,
+            [factory.createIdArgument(hash), factory.createLambdaWrapper(node)]
+        )
+    }
 }
diff --git a/arkoala-arkts/memo-plugin/src/utils.ts b/arkoala-arkts/memo-plugin/src/utils.ts
index 96d5befb48f981233d12be01c7f02ae1d3941695..d936ba040ba8e9e209102689571149bbaa204448 100644
--- a/arkoala-arkts/memo-plugin/src/utils.ts
+++ b/arkoala-arkts/memo-plugin/src/utils.ts
@@ -20,6 +20,7 @@ export enum RuntimeNames {
     __CONTEXT = "__context",
     __ID = "__id",
     ANNOTATION = "memo",
+    COMPUTE = "compute",
     CONTEXT = "__memo_context",
     CONTEXT_TYPE = "__memo_context_type",
     CONTEXT_TYPE_DEFAULT_IMPORT = "@koalaui/runtime",